class_docstr: Adam optimizer with ZeRO algorithm
  
  Child class of Apex DistributedFusedAdam, with optimizations for
  NeMo-Megatron.
  
  Arguments:
      params (iterable): iterable of parameters to optimize or dicts
          defining parameter groups.
      disable_distributed_parameters (bool, optional): use standard
          data-parallel communication instead of ZeRO.
          (default: False)
      **kwargs: keyword arguments to pass to Apex
          DistributedFusedAdam.
code: def sharded_state_dict(self, model_sharded_state_dict):
      optimizer_state_dict = self.state_dict()
  
      id_to_sharded_param_map = get_param_id_to_sharded_param_map(
          model_sharded_state_dict=model_sharded_state_dict,
          optim_params_iter=self.parameters(),
      )
      # Convert state
      step = optimizer_state_dict["state"].pop("step")
      state_dict_format = optimizer_state_dict.pop("format", None)
      optim_state_to_sharding_state(
          optimizer_state_dict, id_to_sharded_param_map
      )
      optimizer_state_dict["state"]["step"] = step
      if state_dict_format is not None:
          optimizer_state_dict["format"] = state_dict_format
  
      def rename_fp32_params(x):
          if isinstance(x, ShardedTensor) and x.key.startswith(
              "optimizer.state.param"
          ):
              x.key = x.key.replace(
                  "optimizer.state.param", "optimizer.state.fp32_param"
              )
          return x
  
      dict_list_map_inplace(rename_fp32_params, optimizer_state_dict)
  
      return optimizer_state_dict
  
imports: import collections
  from typing import Callable, Dict, Iterable, Optional, Union
  import torch
  from apex.contrib.optimizers.distributed_fused_adam import DistributedFusedAdam, _disable_pre_forward_hook, _multi_tensor_copy
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.dict_utils import dict_list_map_inplace
  from megatron.core.dist_checkpointing.mapping import ShardedTensor
  from megatron.core.dist_checkpointing.optimizer import get_param_id_to_sharded_param_map, optim_state_to_sharding_state
  from nemo.utils import str_to_dtype
label_desc: Use image node_id 13 for calling the MegatronDistributedFusedAdam obj's underlying member method code with example usage: obj.sharded_state_dict(model_sharded_state_dict) and returns: optimizer_state_dict, x