code: def _convert_state_with_flat_params(
      all_optim_state_keys: List[_OptimStateKey],
      optim_state_key_to_param_key: Dict[
          _OptimStateKey, Union[int, str]
      ],
      fqn_to_fsdp_param_info: Dict[str, FSDPParamInfo],
      optim_state_dict: Dict[Union[str, int], Any],
      to_save: bool,
      shard_state: bool,
      cpu_offload: bool = True,
  ) -> Dict[str, Any]:
      fsdp_osd_state: Dict[str, Any] = {}
      # Iterate in rank 0's flat 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)
  
          assert param_key is not None, (
              "If use_orig_params is False, we must be able to find the "
              f"corresponding param id. {optim_state_key} {param_key}"
          )
  
          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]
              unflat_state = _unflatten_optim_state(
                  fsdp_param_info,
                  optim_state_dict[param_key],
                  to_save,
                  shard_state,
                  cpu_offload,
              )
              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[param_key]
              )
              if cpu_offload:
                  for state_name, value in sorted_items(
                      fsdp_osd_state[unflat_param_name]
                  ):
                      if not torch.is_tensor(value):
                          continue
                      fsdp_osd_state[unflat_param_name][
                          state_name
                      ] = value.cpu()
  
      return fsdp_osd_state
  
imports: import copy
  import functools
  import logging
  import warnings
  from contextlib import ExitStack
  from dataclasses import dataclass, field
  from typing import Any, cast, Dict, Iterable, Iterator, List, NamedTuple, no_type_check, Optional, Sequence, Set, Tuple, Union
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed._state_dict_utils import _gather_state_dict
  from torch.distributed._tensor import DTensor, Replicate
  from torch.distributed.distributed_c10d import _get_pg_default_device
  from torch.distributed.fsdp._common_utils import _apply_to_modules, _FSDPState, _get_module_fsdp_state_if_fully_sharded_module, _get_param_to_fqns, _module_handle, _named_parameters_with_duplicates, clean_tensor_name
  from torch.distributed.fsdp._debug_utils import SimpleProfiler
  from torch.distributed.fsdp._flat_param import FlatParameter, FlatParamHandle
  from torch.distributed.fsdp._fsdp_extensions import _ext_chunk_dtensor, _ext_chunk_tensor
  from torch.distributed.fsdp._runtime_utils import _lazy_init, _reset_flat_param_grad_info_if_needed
  from torch.distributed.fsdp.api import ShardingStrategy, StateDictSettings, StateDictType
  from torch.utils._pytree import tree_map_only
label_desc: Use image node_id 30 for calling a global function with example usage: _convert_state_with_flat_params(all_optim_state_keys, optim_state_key_to_param_key, fqn_to_fsdp_param_info, optim_state_dict, to_save, shard_state, cpu_offload) and returns: fsdp_osd_state