def _broadcast_pos_dim_tensor_states(
    processed_optim_state_dict: Dict[str, Any],
    flat_optim_state_dict: Optional[Dict[str, Any]],
    rank: int,
    world_size: int,
    group,
    broadcast_device: torch.device,
) -> Dict[str, Any]:
    """
    Takes ``processed_optim_state_dict``, which has metadata in place of
    positive-dimension tensor states, and broadcasts those tensor states from
    rank 0 to all ranks. For tensor states corresponding to FSDP parameters,
    rank 0 shards the tensor and broadcasts shard-by-shard, and for tensor
    states corresponding to non-FSDP parameters, rank 0 broadcasts the full
    tensor.

    Args:
        processed_optim_state_dict (Dict[str, Any]): The flattened optimizer
            state dict with positive-dimension tensor states replaced with
            metadata; this should be returned by
            :meth:`_process_pos_dim_tensor_state` and non-empty on all ranks.
        flat_optim_state_dict (Optional[Dict[str, Any]]): The flattened
            unsharded optimizer state dict with the actual positive-dimension
            tensor states if on rank 0; ignored on nonzero ranks.

    Returns:
        Dict[str, Any]: The optimizer state dict with the positive-dimension
        tensor state correctly populated via ``broadcast()`` s from rank 0.
    """
    assert (
        rank != 0 or flat_optim_state_dict is not None
    ), "Expects rank 0 to pass in the flattened optimizer state dict"
    no_tensor_osd = processed_optim_state_dict  # alias
    flat_osd = flat_optim_state_dict  # alias
    for key, param_state in no_tensor_osd["state"].items():
        for state_name, value in sorted_items(param_state):
            is_pos_dim_tensor_state = isinstance(value, _PosDimTensorInfo)
            if not is_pos_dim_tensor_state:
                continue
            if rank == 0:
                assert flat_osd is not None
                unsharded_tensor = flat_osd["state"][key][state_name]
            else:
                unsharded_tensor = None
            shape, dtype = value.shape, value.dtype
            if key.is_fsdp_managed:  # FSDP parameter
                _broadcast_sharded_pos_dim_tensor_state(
                    unsharded_tensor,
                    param_state,
                    state_name,
                    shape,
                    dtype,
                    broadcast_device,
                    rank,
                    world_size,
                    group,
                )  # modify `param_state` destructively
            else:  # non-FSDP parameter
                _broadcast_unsharded_pos_dim_tensor_state(
                    unsharded_tensor,
                    param_state,
                    state_name,
                    shape,
                    dtype,
                    broadcast_device,
                    rank,
                    group,
                )  # modify `param_state` destructively
    return no_tensor_osd
Takes ``processed_optim_state_dict``, which has metadata in place of
positive-dimension tensor states, and broadcasts those tensor states from
rank 0 to all ranks. For tensor states corresponding to FSDP parameters,
rank 0 shards the tensor and broadcasts shard-by-shard, and for tensor
states corresponding to non-FSDP parameters, rank 0 broadcasts the full
tensor.

Args:
processed_optim_state_dict (Dict[str, Any]): The flattened optimizer
state dict with positive-dimension tensor states replaced with
metadata; this should be returned by
:meth:`_process_pos_dim_tensor_state` and non-empty on all ranks.
flat_optim_state_dict (Optional[Dict[str, Any]]): The flattened
unsharded optimizer state dict with the actual positive-dimension
tensor states if on rank 0; ignored on nonzero ranks.

Returns:
Dict[str, Any]: The optimizer state dict with the positive-dimension
tensor state correctly populated via ``broadcast()`` s from rank 0.
_optim_utils.py