def _flatten_optim_state_dict(
    optim_state_dict: Dict[str, Any],
    model: nn.Module,
    shard_state: bool,
    use_orig_params: bool = False,
    optim: Optional[torch.optim.Optimizer] = None,
) -> Dict[str, Any]:
    """
    Flattens the full optimizer state dict, still keying by unflattened
    parameter names. If ``shard_state=True``, then FSDP-managed
    ``FlatParameter`` 's optimizer states are sharded, and otherwise, they are
    kept unsharded.

    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.

    Note that ``_flatten_tensor_optim_state`` does not need ``optim`` to
    flatten/shard the state. However, NamedOptimizer and KeyedOptimizer require
    all the states even if the corresponding parameters are empty. To this end,
    ``optim`` will be used to to get the initial state of the empty parameters.
    ``optim`` should only be non-None if the ``optim` is KeyedOptimizer or
    NamedOptimizer.

    Returns:
        Dict[str, Any]: The flattened optimizer state dict.
    """
    unflat_osd = optim_state_dict
    if "state" not in unflat_osd or "param_groups" not in unflat_osd:
        raise ValueError(
            '`optim_state_dict` must have the keys "state" and '
            '"param_groups" to be a valid optimizer state dict'
        )
    param_to_fqns = _get_param_to_fqns(model)
    fqn_to_fsdp_param_info = _get_fqn_to_fsdp_param_info(model)

    # Construct the "state" part
    flat_osd_state: Dict[Union[_OptimStateKey, str], Any] = {}
    unflat_osd_state = unflat_osd["state"]
    all_state_keys = set(unflat_osd_state.keys())

    # local_state_dict is used to construct states of empty parameters.
    # This should only be used if is_named_optimizer=True.
    local_state_dict: Dict[str, Any] = {}
    local_state_clean_fqns: Dict[str, str] = {}
    if optim is not None:
        local_state_dict = optim.state_dict()["state"]
        for fqn in local_state_dict.keys():
            clean_fqn = clean_tensor_name(fqn)
            local_state_clean_fqns[clean_fqn] = fqn

    for param, unflat_param_names in param_to_fqns.items():
        fqn = unflat_param_names[0]
        if fqn not in unflat_osd_state:
            continue
        all_state_keys.difference_update(unflat_param_names)
        if fqn in fqn_to_fsdp_param_info:
            fsdp_param_info = fqn_to_fsdp_param_info[fqn]
            if use_orig_params:
                assert (
                    shard_state
                ), "If use_orig_params is True, shard_state must be True."
                flat_state = _shard_orig_param_state(
                    fsdp_param_info,
                    fqn,
                    unflat_osd_state[fqn],
                )
            else:
                flat_state = _flatten_optim_state(
                    fsdp_param_info,
                    unflat_osd_state,
                    unflat_param_names,
                    shard_state,
                )
            key = _OptimStateKey(tuple(unflat_param_names), True)
            # Only include non-empty states since as expected by
            # `torch.optim.Optimizer` s unless the optimizer is KeyedOptimizer
            # or NamedOptimizer.
            if flat_state:
                flat_osd_state[key] = flat_state
            elif optim is not None:  # NamedOptimizer or KeyedOptimizer case.
                assert len(unflat_param_names) == 1
                local_wrapped_fqn = local_state_clean_fqns.get(fqn, "")
                if local_wrapped_fqn:
                    flat_osd_state[key] = copy.deepcopy(
                        local_state_dict[local_wrapped_fqn]
                    )
        else:  # do not flatten non-FSDP parameters' states
            assert len(unflat_param_names) == 1
            key = _OptimStateKey(tuple(unflat_param_names), False)
            flat_osd_state[key] = copy.copy(unflat_osd_state[fqn])

    # Handle user-defined state, states that are not accosiated with parameters.
    for key in all_state_keys:
        flat_osd_state[key] = copy.copy(unflat_osd_state[key])

    # Construct the "param_groups" part -- copy as is since it will be
    # rekeyed later according to the target rank's optimizer
    flat_osd_param_groups = copy.deepcopy(unflat_osd["param_groups"])
    return {"state": flat_osd_state, "param_groups": flat_osd_param_groups}
Flattens the full optimizer state dict, still keying by unflattened
parameter names. If ``shard_state=True``, then FSDP-managed
``FlatParameter`` 's optimizer states are sharded, and otherwise, they are
kept unsharded.

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.

Note that ``_flatten_tensor_optim_state`` does not need ``optim`` to
flatten/shard the state. However, NamedOptimizer and KeyedOptimizer require
all the states even if the corresponding parameters are empty. To this end,
``optim`` will be used to to get the initial state of the empty parameters.
``optim`` should only be non-None if the ``optim` is KeyedOptimizer or
NamedOptimizer.

Returns:
Dict[str, Any]: The flattened optimizer state dict.
_optim_utils.py