def _gather_orig_param_state(
    fsdp_param_info: FSDPParamInfo,
    fqn: str,
    optim_state: Dict[str, Any],
    shard_state: bool,
    group=None,
) -> Dict[str, Any]:
    """
    Gather the optimizer state for the original parameter with the name ``fqn``.
    This API should only be used when ``use_orig_params`` is True.
    """
    fsdp_state = fsdp_param_info.state
    assert (
        fsdp_state._use_orig_params
    ), "_gather_orig_param_state only support use_orig_params=True case"
    flat_param = fsdp_param_info.flat_param
    param_idx = fsdp_param_info.param_indices[fqn]
    if (
        fsdp_state.world_size == 1
        or fsdp_state.sharding_strategy == ShardingStrategy.NO_SHARD
    ):
        return optim_state

    gathered_state = _all_gather_optim_state(fsdp_state, optim_state, group=group)

    # Unflatten state values.
    for state_name, value in list(gathered_state.items()):
        if not torch.is_tensor(value) or value.dim() == 0:
            continue

        value = value[: flat_param._numels[param_idx]].reshape(
            flat_param._shapes[param_idx]
        )
        if shard_state:
            assert fsdp_state.process_group is not None
            value = _ext_chunk_tensor(
                value,
                fsdp_state.rank,
                fsdp_state.world_size,
                torch.cuda.device_count(),
                fsdp_state.process_group,
            )
        value = value.cpu()
        gathered_state[state_name] = value
    return gathered_state
Gather the optimizer state for the original parameter with the name ``fqn``.
This API should only be used when ``use_orig_params`` is True.
_optim_utils.py