def _unflatten_communicated_optim_state(
    fsdp_param_info: FSDPParamInfo,
    state: _ConsolidatedOptimState,
    shard_state: bool,
) -> List[Dict[str, Any]]:
    """
    Unflattens the communicated optimizer state (given by ``tensor_state``,
    ``non_tensor_state``, and ``zero_dim_tensor_state``) for a single flattened
    parameter. This should only be called on the target rank.

    Args:
        fsdp_param_info (FSDPParamInfo): The fsdp state and the target flatten
            parameter.
        state (_ConsolidatedOptimState): Consolidated optimizer state.

    Returns:
        List[Dict[str, Any]]: A :class:`list` holding the entries in the
        "state" part of the optimizer state dict corresponding to the
        unflattened parameters comprising the flattened parameter. The final
        optimizer state dict will need to map these entries using the proper
        unflattened parameter IDs.
    """
    fsdp_state = fsdp_param_info.state
    flat_param = fsdp_param_info.flat_param
    unflat_param_state: List[Dict[str, Any]] = []
    flat_param_views: Dict[str, Iterator] = {}
    num_unflat_params = flat_param._num_params
    tensor_state, zero_dim_tensor_state, non_tensor_state = (
        state.tensor_state,
        state.zero_dim_tensor_state,
        state.non_tensor_state,
    )

    for _ in range(num_unflat_params):
        unflat_state_param = {}
        # Add positive-dimension tensor state: unflatten with views
        for state_name, flat_tensor in sorted_items(tensor_state):
            views_generated = state_name in flat_param_views
            if not views_generated:
                views = FlatParamHandle._get_unflat_views(flat_param, flat_tensor)
                flat_param_views[state_name] = views
            else:
                views = flat_param_views[state_name]
            optim_state: Union[torch.Tensor, ShardedTensor] = next(views)
            if shard_state:
                assert fsdp_state.process_group is not None
                optim_state = _ext_chunk_tensor(
                    optim_state,
                    fsdp_state.rank,
                    fsdp_state.world_size,
                    torch.cuda.device_count(),
                    fsdp_state.process_group,
                )
            unflat_state_param[state_name] = optim_state

        # Add zero-dimension tensor state: take the target rank's value
        for state_name, zero_dim_tensor in sorted_items(zero_dim_tensor_state):
            unflat_state_param[state_name] = zero_dim_tensor
        # Add non-tensor state: take the target rank's value
        for state_name, non_tensor in sorted_items(non_tensor_state):
            unflat_state_param[state_name] = non_tensor
        unflat_param_state.append(unflat_state_param)
    return unflat_param_state
Unflattens the communicated optimizer state (given by ``tensor_state``,
``non_tensor_state``, and ``zero_dim_tensor_state``) for a single flattened
parameter. This should only be called on the target rank.

Args:
fsdp_param_info (FSDPParamInfo): The fsdp state and the target flatten
parameter.
state (_ConsolidatedOptimState): Consolidated optimizer state.

Returns:
List[Dict[str, Any]]: A :class:`list` holding the entries in the
"state" part of the optimizer state dict corresponding to the
unflattened parameters comprising the flattened parameter. The final
optimizer state dict will need to map these entries using the proper
unflattened parameter IDs.
_optim_utils.py