def _process_pos_dim_tensor_state(
    flat_optim_state_dict: Dict[str, Any],
    world_size: int,
) -> Dict[str, Any]:
    """
    Processes positive-dimension tensor states in ``flat_optim_state_dict`` by
    replacing them with metadata. This is done so the processed optimizer state
    dict can be broadcast from rank 0 to all ranks without copying those tensor
    states, and thus, this is meant to only be called on rank 0.

    Args:
        flat_optim_state_dict (Dict[str, Any]): Flattened optimizer state dict
            with the positive-dimension tensor states unsharded.

    Returns:
        Dict[str, Any]: The flattened optimizer state dict with positive-
        dimension tensor states replaced by metadata.
    """
    flat_osd = flat_optim_state_dict  # alias
    no_tensor_osd: Dict[str, Any] = {"state": {}}
    for key, param_state in flat_osd["state"].items():
        no_tensor_osd["state"][key] = {}
        for state_name, value in sorted_items(param_state):
            is_pos_dim_tensor_state = torch.is_tensor(value) and value.dim() > 0
            if not is_pos_dim_tensor_state:
                no_tensor_osd["state"][key][state_name] = value
                continue
            if key.is_fsdp_managed:  # FSDP parameter
                sharded_size = FlatParamHandle._get_sharded_size(
                    value, rank=0, world_size=world_size
                )
                assert len(sharded_size) == 1, f"{sharded_size}"
                info = _PosDimTensorInfo(sharded_size, value.dtype)
            else:  # non-FSDP parameter
                info = _PosDimTensorInfo(value.shape, value.dtype)
            no_tensor_osd["state"][key][state_name] = info
    no_tensor_osd["param_groups"] = flat_osd["param_groups"]
    return no_tensor_osd
Processes positive-dimension tensor states in ``flat_optim_state_dict`` by
replacing them with metadata. This is done so the processed optimizer state
dict can be broadcast from rank 0 to all ranks without copying those tensor
states, and thus, this is meant to only be called on rank 0.

Args:
flat_optim_state_dict (Dict[str, Any]): Flattened optimizer state dict
with the positive-dimension tensor states unsharded.

Returns:
Dict[str, Any]: The flattened optimizer state dict with positive-
dimension tensor states replaced by metadata.
_optim_utils.py