code: def _unflatten_optim_state(
      combined_state: Dict[int, Dict],
      instance_list: List["FullyShardedDataParallel"],
      world_pad_info: List[List[List[int]]],
      singleton_state: Dict[int, Dict],
  ) -> Tuple[Dict[int, Dict], Dict[int, int]]:
      """Convert optimizer state for flattened parameters into original, unflattened ones.
  
      Args:
          combined_state: all-gathered state with tensors
          instance_list: list of FSDP wrapper object instances
          world_pad_info: [param_id][fsdp_instance_id][bytes_padded_per_rank]
          singleton_state: all-gathered dimensionless tensors
  
      Returns:
          state: unflattened state dict
          idx_mapping: a mapping from global ID to local ID
      """
      # local ids are the keys in the current state (combined_state), (usually fewer)
      # global ids will be the keys in the unflattened state
      next_global_id = 0  # gets incremented
      pad_info = {
          id: [s[id][0] for s in world_pad_info]
          for id in combined_state
      }
      local_ids = [id for id in sorted(combined_state.keys())]
  
      # non_tensor_state refers to entries in sd[state][param_id] that are not tensors, like "step".
      # we check that these are identical across workers and then take the first
      non_tensor_state = {
          id: _extract_non_tensor_state(combined_state, id)
          for id in combined_state
      }
  
      # Local corresponds to flattened, global corresponds to unflattened.
      # Casting needed only for mypy.
      num_global_params: List[int] = []
      for m in instance_list:
          if m.flatten_parameters:
              num_flatten = cast(int, m.num_params_managed)
              num_global_params.append(num_flatten)
          else:
              num_global_params.append(len(m.non_shared_params()))
      global_to_local_id = {}
      for local_id, num_unflat in enumerate(num_global_params):
          for _ in range(num_unflat):
              # Some params could be unused, which means the optimizer
              # hasn't created their state. Therefore, `local_id` obtained
              # by enumerating the params above could be out of the range
              # of keys in `combined_state` above. Here is an example:
              #
              #    global    local    notes
              #    0         0        FC1's weight, first flat buffer
              #    1         0        FC1's bias, first flat buffer
              #    2         None     FC2's weight, no flat state
              #    3         None     FC2's bias, no flat state
              #    4         2        FC3's weight, second flat buffer (but with id 2)
              #    5         2        FC3's bias, second flat buffer (but with id 2)
              global_to_local_id[next_global_id] = (
                  local_id if local_id in local_ids else None
              )
              next_global_id += 1
      if not combined_state:
          return {}, global_to_local_id
  
      # copy non tensor state (like the "step" count) to all global entries
      unflat_state = {
          i: copy.deepcopy(non_tensor_state[0])
          for i in range(sum(num_global_params))
      }
  
      # remove the global entries that don't have optim state because pytorch
      # optimizer's state_dict() function returns a state_dict without the missing
      # param, so we shouldn't have things like "1:{}" for missing params.
      for g, l in global_to_local_id.items():
          if l is None:
              del unflat_state[g]
  
      if non_tensor_state[0].keys() == combined_state[0].keys():
          # Early return if there is no tensors in the state dict.
          return unflat_state, global_to_local_id
  
      local_to_global: Dict[int, List] = {i: [] for i in local_ids}
      for g, l in global_to_local_id.items():
          if l is not None:
              local_to_global[l].append(g)
      # loop over parameters in state.
      # Tensor state will be padded, concatenated, and restored to original shape with FlattenParamsWrapper.get_views
      # get_views returns multiple tensors, each of which is a new parameter with a new "global" id.
      for local_id in local_ids:
          # undo the work of shard_parameters
          for k, v in combined_state[local_id].items():
              if k in non_tensor_state[local_id]:
                  continue
              assert isinstance(v, list), f"got {k}: {v} for {local_id}"
              v_unpad = [
                  t[:-np] if np > 0 else t
                  for t, np in zip(v, pad_info[local_id])
              ]
              flat_buffer = torch.cat(v_unpad)
              if instance_list[local_id].flatten_parameters:
                  # Unflatten. Casting needed only for mypy.
                  param_views: Iterator = cast(
                      FlattenParamsWrapper, instance_list[local_id]
                  ).get_param_views([flat_buffer])
                  for global_id, param_view in zip(
                      sorted(local_to_global[local_id]), param_views
                  ):
                      assert (
                          k not in unflat_state[global_id]
                      ), f"already added {k} to {global_id} {local_id}"
                      unflat_state[global_id][k] = param_view
              else:
                  # Copy non-flatten state directly.
                  assert (
                      len(local_to_global[local_id]) == 1
                  ), "Only support a single non-flatten parameter"
                  global_id = local_to_global[local_id][0]
                  unflat_state[global_id][k] = flat_buffer
              unflat_state[global_id].update(singleton_state[local_id])
  
      return unflat_state, global_to_local_id
  
file_path: fairscale/fairscale/nn/data_parallel/fsdp_optim_utils.py
imports: import copy
  from itertools import groupby
  from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Tuple, cast
  import torch
  from fairscale.nn.misc import FlattenParamsWrapper
label_desc: Use image node_id 4 for calling a global function with example usage: _unflatten_optim_state(combined_state, instance_list, world_pad_info, singleton_state) and returns: unflat_state,  global_to_local_id, dict,  global_to_local_id, unflat_state,  global_to_local_id