code: def get_zero_optimizer_state_dict_local(optimizer, global_rank):
      optimizer._check_overlap_initialized()
  
      # Sync the exposed `param_groups` attributes to the local optimizer in
      # case they have been updated
      optimizer._sync_param_groups(
          optimizer.param_groups, optimizer.optim.param_groups
      )
  
      local_state_dict = optimizer.optim.state_dict()
      state_dict = super(
          ZeroRedundancyOptimizer, optimizer
      ).state_dict()
  
      # Update the global optimizer state with local state information,
      # factoring in the translation from local to global indexing
      rank = global_rank
      # TODO: recursive copy to device
      local_param_groups = local_state_dict["param_groups"]
      global_param_groups = optimizer._partition_parameters()[rank]
      assert len(local_param_groups) == len(
          global_param_groups
      ), "Mismatch between number of local and global parameter groups"
  
      for local_param_group, global_param_group in zip(
          local_param_groups, global_param_groups
      ):
          # `local_param_group` stores local indices, while
          # `global_param_group` stores the tensors directly
          local_param_indices = local_param_group["params"]
          global_params = global_param_group["params"]
  
          assert len(local_param_indices) == len(
              global_params
          ), "Mismatch between number of local and global parameters in parameter group"
          for local_param_index, global_param in zip(
              local_param_indices, global_params
          ):
              # Update the global parameter state, if any
              if local_param_index in local_state_dict["state"]:
                  global_param_index = optimizer._param_to_index[
                      global_param
                  ]
                  state_dict["state"][
                      global_param_index
                  ] = local_state_dict["state"][local_param_index]
  
      # Sort the parameters in the state
      state_dict["state"] = dict(sorted(state_dict["state"].items()))
      return state_dict
  
file_path: flash-attention/training/src/utils/ddp_zero1.py
imports: from typing import Any, Callable, Dict, List, Optional, Union
  from pathlib import Path
  import torch
  from torch.optim.optimizer import Optimizer
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from pytorch_lightning.strategies.ddp import DDPStrategy
  from pytorch_lightning.core.optimizer import LightningOptimizer
label_desc: Use image node_id 1 for calling a global function with example usage: get_zero_optimizer_state_dict_local(optimizer, global_rank) and returns: state_dict