def _optim_state_dict(
    model: nn.Module,
    optim: torch.optim.Optimizer,
    optim_state_dict: Dict[str, Any],
    optim_input: Optional[
        Union[
            List[Dict[str, Any]],
            Iterable[nn.Parameter],
        ]
    ],
    rank0_only: bool,
    shard_state: bool,
    group: Optional[dist.ProcessGroup],
    using_optim_input: bool,
    use_orig_params: bool = False,
) -> Dict[str, Any]:
    """
    Consolidates the optimizer state and returns it as a :class:`dict`
    following the convention of :meth:`torch.optim.Optimizer.state_dict`,
    i.e. with keys ``"state"`` and ``"param_groups"``.
    The flattened parameters in ``FSDP`` modules contained in ``model``
    are mapped back to their unflattened parameters.

    Parameter keys are not well-defined. For a regular optimizer, the optimizer
    state_dict contains a mapping from parameter IDs to parameter states.
    Parameter IDs are the order of parameters in ``optim.param_groups()`` across
    all the groups. This API also allows user to pass ``optim_input`` for the
    mapping between parameters and parameter IDs. Using ``optim_input`` is being
    deprecated.

    If the optimizer is a ``NamedOptimizer``, the optimizer state_dict does not
    contain parameter IDs mapping but a mapping from parameter FQNs to parameter
    states. This API finds the mapping from FQNs to parameters if the optimizer
    is a ``NamedOptimizer``.

    If ``use_orig_params`` is True, each rank will have all FSDP-managed
    parameters but some of these parameters may be empty due to the sharding.
    For a regular optim.Optimizer, states for those empty parameters will
    not be initialized. So, when aggregating the FQNs across ranks, no assert
    will be raised on a rank even if it does not have all the states -- it is
    valid and FSDP know how to aggregate them. However, FSDP has to ignore
    handling those parameters that are not managed by FSDP and do not exist on
    the local rank -- it is managed by other parallelism and FSDP does not
    know ho to handle/aggregate them.

    Args:
        model (nn.Module): Root module (which may or may not be a
            :class:`FullyShardedDataParallel` instance) whose parameters
            were passed into the optimizer ``optim``.
        optim (torch.optim.Optimizer): Optimizer for ``model`` 's
            parameters.
        rank0_only (bool): If ``True``, saves the populated :class:`dict`
            only on rank 0; if ``False``, saves it on all ranks. (Default:
            ``True``)
        shard_state (bool): If ``True``, shard and distribute all
            non-zero-dimension states.

    Returns:
        Dict[str, Any]: A :class:`dict` containing the optimizer state for
        ``model`` 's original unflattened parameters and including keys
        "state" and "param_groups" following the convention of
        :meth:`torch.optim.Optimizer.state_dict`. If ``rank0_only=False``,
        then nonzero ranks return an empty :class:`dict`.
    """
    _clear_grads_if_needed(traversal_utils._get_fsdp_handles(model))
    to_save = not rank0_only or (dist.get_rank(group) == 0 or shard_state)
    fsdp_osd: Dict[str, Any] = {"state": {}, "param_groups": []} if to_save else {}
    fsdp_osd_state: Dict[str, Any] = fsdp_osd["state"] if to_save else {}
    param_to_fqns = _get_param_to_fqns(model)
    flat_param_to_fqn = _get_flat_param_to_fqn(model)
    is_named_optimizer = _is_named_optimizer(optim_state_dict)

    param_key_to_param = cast(
        Dict[Union[int, str], nn.Parameter],
        (
            _get_param_id_to_param_from_optim_input(model, optim_input)
            if using_optim_input
            else _get_param_key_to_param(
                optim, model, is_named_optimizer, param_to_fqns, flat_param_to_fqn
            )
        ),
    )
    fqn_to_fsdp_param_info = _get_fqn_to_fsdp_param_info(model)

    all_optim_state_keys, optim_state_key_to_param_key = _map_param_key_to_optim_keys(
        optim_state_dict,
        group,
        param_key_to_param,
        param_to_fqns,
        fqn_to_fsdp_param_info,
        merge_keys=use_orig_params,
    )

    # Iterate in rank 0's flattened parameter ID order to ensure aligned
    # all-gathers across ranks
    for optim_state_key in all_optim_state_keys:
        param_key: Union[str, int, None] = optim_state_key_to_param_key.get(
            optim_state_key, None
        )

        if param_key is None:
            assert use_orig_params, (
                "If use_orig_params is False, we must be able to find the "
                f"corresponding param id. {optim_state_key} {param_key}"
            )
            if not optim_state_key.is_fsdp_managed:
                continue

        if optim_state_key.is_fsdp_managed:
            # If there are multiple unflat_param_names (not use_orig_params),
            # they share the same FSDPParamInfo. So the first unflat_param_name
            # is sufficient to fetch the FSDPParamInfo.
            fqn = optim_state_key.unflat_param_names[0]
            fsdp_param_info = fqn_to_fsdp_param_info[fqn]
            if use_orig_params:
                state = (
                    {} if param_key is None else optim_state_dict["state"][param_key]
                )
                unflat_state = [
                    _gather_orig_param_state(
                        fsdp_param_info, fqn, state, shard_state, group
                    )
                ]
            else:
                unflat_state = _unflatten_optim_state(
                    fsdp_param_info,
                    optim_state_dict["state"][param_key],
                    to_save,
                    shard_state,
                )
            if to_save:
                assert len(unflat_state) == len(optim_state_key.unflat_param_names)
                for unflat_param_name, unflat_param_state in zip(
                    optim_state_key.unflat_param_names,
                    unflat_state,
                ):
                    fsdp_osd_state[unflat_param_name] = unflat_param_state
        elif to_save:
            assert len(optim_state_key.unflat_param_names) == 1
            unflat_param_name = optim_state_key.unflat_param_names[0]
            fsdp_osd_state[unflat_param_name] = copy.copy(
                optim_state_dict["state"][param_key]
            )
            for state_name, value in sorted_items(fsdp_osd_state[unflat_param_name]):
                if torch.is_tensor(value):
                    fsdp_osd_state[unflat_param_name][state_name] = value.cpu()

    if to_save:
        flat_param_fqns = set(flat_param_to_fqn.values())
        for key, value in optim_state_dict["state"].items():
            if key in fsdp_osd_state:
                continue
            if key in flat_param_fqns:
                continue
            if key in param_key_to_param:
                continue
            # This key is not recognized by FSDP. It may be a user-defined state
            # or some parameters state that FSDP is unable to map from
            # ``optim.param_groups``.
            warnings.warn(
                f"Found a optim state, {key}, that FSDP cannot process. FSDP "
                "will directly copy everything to the returned state_dict. In "
                "most cases, this is a user-defined state that is not "
                "associated with any particular parameter. Another possible "
                "case is this state is managed by DMP. Otherwise, there may "
                " be a mismatched assumption of optim_state_dict of this mode."
            )
            fsdp_osd_state[key] = value

        fsdp_osd["param_groups"] = _unflatten_param_groups(
            optim_state_dict, param_key_to_param, param_to_fqns
        )

    return fsdp_osd
Consolidates the optimizer state and returns it as a :class:`dict`
following the convention of :meth:`torch.optim.Optimizer.state_dict`,
i.e. with keys ``"state"`` and ``"param_groups"``.
The flattened parameters in ``FSDP`` modules contained in ``model``
are mapped back to their unflattened parameters.

Parameter keys are not well-defined. For a regular optimizer, the optimizer
state_dict contains a mapping from parameter IDs to parameter states.
Parameter IDs are the order of parameters in ``optim.param_groups()`` across
all the groups. This API also allows user to pass ``optim_input`` for the
mapping between parameters and parameter IDs. Using ``optim_input`` is being
deprecated.

If the optimizer is a ``NamedOptimizer``, the optimizer state_dict does not
contain parameter IDs mapping but a mapping from parameter FQNs to parameter
states. This API finds the mapping from FQNs to parameters if the optimizer
is a ``NamedOptimizer``.

If ``use_orig_params`` is True, each rank will have all FSDP-managed
parameters but some of these parameters may be empty due to the sharding.
For a regular optim.Optimizer, states for those empty parameters will
not be initialized. So, when aggregating the FQNs across ranks, no assert
will be raised on a rank even if it does not have all the states -- it is
valid and FSDP know how to aggregate them. However, FSDP has to ignore
handling those parameters that are not managed by FSDP and do not exist on
the local rank -- it is managed by other parallelism and FSDP does not
know ho to handle/aggregate them.

Args:
model (nn.Module): Root module (which may or may not be a
:class:`FullyShardedDataParallel` instance) whose parameters
were passed into the optimizer ``optim``.
optim (torch.optim.Optimizer): Optimizer for ``model`` 's
parameters.
rank0_only (bool): If ``True``, saves the populated :class:`dict`
only on rank 0; if ``False``, saves it on all ranks. (Default:
``True``)
shard_state (bool): If ``True``, shard and distribute all
non-zero-dimension states.

Returns:
Dict[str, Any]: A :class:`dict` containing the optimizer state for
``model`` 's original unflattened parameters and including keys
"state" and "param_groups" following the convention of
:meth:`torch.optim.Optimizer.state_dict`. If ``rank0_only=False``,
then nonzero ranks return an empty :class:`dict`.
_optim_utils.py