def _broadcast_unsharded_pos_dim_tensor_state(
    unsharded_tensor: Optional[torch.Tensor],
    param_state: Dict[str, Any],
    state_name: str,
    shape: torch.Size,
    dtype: torch.dtype,
    broadcast_device: torch.device,
    rank: int,
    group,
) -> None:
    """
    Broadcasts positive-dimension tensor state for the state ``state_name``
    corresponding to an unsharded non-FSDP parameter from rank 0 to all ranks.
    This modifies ``param_state`` destructively.

    Args:
        unsharded_tensor (Optional[torch.Tensor]): Unsharded tensor to
            broadcast if on rank 0; ignored otherwise.
    """
    if rank == 0:
        assert (
            unsharded_tensor is not None
        ), "Expects rank 0 to pass in the unsharded tensor"
        assert (
            shape == unsharded_tensor.shape
        ), f"Shape mismatch: {shape} {unsharded_tensor.shape}"
        assert (
            dtype == unsharded_tensor.dtype
        ), f"dtype mismatch: {dtype} {unsharded_tensor.dtype}"
        unsharded_tensor = unsharded_tensor.to(broadcast_device)
    else:
        unsharded_tensor = torch.zeros(
            shape,
            requires_grad=False,
            dtype=dtype,
            device=broadcast_device,
        )
    dist.broadcast(unsharded_tensor, src=0, group=group)
    # Keep the tensor on the broadcast device, which is typically GPU
    param_state[state_name] = unsharded_tensor
Broadcasts positive-dimension tensor state for the state ``state_name``
corresponding to an unsharded non-FSDP parameter from rank 0 to all ranks.
This modifies ``param_state`` destructively.

Args:
unsharded_tensor (Optional[torch.Tensor]): Unsharded tensor to
broadcast if on rank 0; ignored otherwise.
_optim_utils.py