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 __init__(
      self,
      params: Union[Iterable[torch.nn.Parameter], Iterable[dict]],
      disable_distributed_parameters: bool = False,
      **kwargs,
  ):
      # Initialize process groups
      if (
          "process_group" not in kwargs
          and not parallel_state.is_unitialized()
      ):
          kwargs[
              "process_group"
          ] = parallel_state.get_data_parallel_group(
              with_context_parallel=True
          )
      if disable_distributed_parameters:
          world_size = torch.distributed.get_world_size()
          rank = torch.distributed.get_rank()
          self_groups = [
              torch.distributed.new_group(ranks=[i])
              for i in range(world_size)
          ]
          kwargs["distributed_process_group"] = self_groups[rank]
          kwargs["redundant_process_group"] = kwargs["process_group"]
  
      # Make sure dtypes are in right type
      for keyword in ("dtype", "grad_sync_dtype", "param_sync_dtype"):
          if keyword in kwargs:
              kwargs[keyword] = str_to_dtype(kwargs[keyword])
  
      # Make sure params are in consistent format (list of param group dicts)
      param_groups = list(params)
      assert param_groups
      if not isinstance(param_groups[0], dict):
          param_groups = [{"params": param_groups}]
  
      # Construct distributed optimizer
      super().__init__(param_groups, **kwargs)
  
      # Initialize weights that require FP32 grads
      if (
          self.dtype != torch.float32
          or self.grad_sync_dtype != torch.float32
      ):
          fp32_params = []
          for param_group in param_groups:
              fp32_params.extend(
                  filter(
                      lambda param: getattr(
                          param, "_with_fp32_optimizer", False
                      ),
                      param_group["params"],
                  )
              )
          if fp32_params:
              assert (
                  self.dtype == torch.float32
              ), f"Param requires FP32 state but optimizer is initialized with {self.dtype}"
              self.init_params_bucket(
                  fp32_params,
                  grad_sync_dtype=torch.float32,
              )
  
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 1 to create a new MegatronDistributedFusedAdam object from inherited base classes: DistributedFusedAdam with example: obj = MegatronDistributedFusedAdam(params, disable_distributed_parameters)