def _all_gather_optim_state(
    fsdp_state: _FSDPState,
    optim_state: Dict[str, Any],
    group=None,
) -> Dict[str, Any]:
    """
    All-gathering state from all the ranks. This API is slow as it uses
    ``all_gather_object``. However, optim state_dict is not in the critical path.
    We can fuse the communication across differnt state if the performance
    becomes a problem.
    """
    # Allgather the scalar tensor state, non-tensor states and tensors metadata.
    processed_state = StateInfo({}, {}, {})
    for state_name, value in sorted_items(optim_state):
        if torch.is_tensor(value):
            if value.dim() == 0:
                # Ensure that `step` is on CPU.
                processed_state.scalar_tensors[state_name] = value.cpu()
            else:
                processed_state.tensors[state_name] = _PosDimTensorInfo(
                    value.shape, value.dtype
                )
        else:
            processed_state.non_tensors = value
    object_list: List[StateInfo] = [
        processed_state for _ in range(fsdp_state.world_size)
    ]
    dist.all_gather_object(object_list, processed_state, group=group)

    # Convert the gathered, pre-proccessed state of each rank to the original one.
    gathered_state: Dict[str, Any] = {}

    all_tensor_states = sorted(
        {n for state in object_list for n in state.tensors.keys()}
    )
    empty_ranks: Set[int] = set()
    for name in all_tensor_states:
        numels = []
        dtype = torch.float
        _empty_ranks: Set[int] = set()
        for rank, object_state in enumerate(object_list):
            numels.append(0)
            info = object_state.tensors.get(name, None)
            if info is not None:
                numels[-1] = info.shape.numel()
                dtype = info.dtype
            if numels[-1] == 0:
                _empty_ranks.add(rank)

        empty_func = functools.partial(
            torch.empty, dtype=dtype, device=fsdp_state.compute_device
        )
        if empty_ranks:
            assert empty_ranks == _empty_ranks
        empty_ranks = _empty_ranks
        local_state = optim_state.get(name, empty_func(0))
        local_state = local_state.to(fsdp_state.compute_device)
        tensors = [
            empty_func(numel) if rank != fsdp_state.rank else local_state
            for rank, numel in enumerate(numels)
        ]
        work = dist.all_gather(
            tensors, local_state, group=fsdp_state.process_group, async_op=True
        )
        gathered_state[name] = AllGatherInfo(tensors, numels, work)

    for rank, object_state in enumerate(object_list):
        if rank in empty_ranks:
            continue
        for name, non_tensor_value in object_state.non_tensors.items():
            curr_non_tensor_value = gathered_state.get(name, None)
            assert (
                curr_non_tensor_value is None
                or curr_non_tensor_value == non_tensor_value
            ), f"Different ranks have different values for {name}."
            gathered_state[name] = non_tensor_value

        for name, scalar_tensor_value in object_state.scalar_tensors.items():
            curr_scalar_tensor_value = gathered_state.get(name, None)
            assert curr_scalar_tensor_value is None or torch.equal(
                scalar_tensor_value, curr_scalar_tensor_value
            ), f"Different ranks have different values for {name}."
            gathered_state[name] = scalar_tensor_value

    for name, value in list(gathered_state.items()):
        if not isinstance(value, AllGatherInfo):
            continue
        assert value.work is not None
        value.work.wait()
        gathered_state[name] = torch.cat(
            [
                rank_tensor[:rank_numel]
                for rank_tensor, rank_numel in zip(value.tensors, value.numels)
                if rank_numel > 0
            ]
        )

    return gathered_state
All-gathering state from all the ranks. This API is slow as it uses
``all_gather_object``. However, optim state_dict is not in the critical path.
We can fuse the communication across differnt state if the performance
becomes a problem.
_optim_utils.py