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 grad_norm(
      self,
      parameters: Optional[Iterable[torch.nn.Parameter]] = None,
      norm_type: float = 2.0,
      force: bool = False,
  ) -> torch.Tensor:
      assert norm_type == 2
  
      if parameters is not None:
          # Make sure we can access iterable multiple times
          parameters = list(parameters)
  
      # Compute grad norm
      if force or self._grad_norm is None:
          # Compute norm of local gradients for distributed optimizer
          grad_norm_sq = self._local_grad_norm(
              parameters=parameters, norm_type=norm_type
          )
          if self.redundant_size > 1:
              grad_norm_sq /= self.redundant_size
  
          # Sum over all procs to get grad norm
          torch.distributed.all_reduce(
              grad_norm_sq,
              op=torch.distributed.ReduceOp.SUM,
          )
          self._grad_norm = grad_norm_sq.sqrt()
  
      # Use cached grad norm
      return super().grad_norm()
  
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 10 for calling the MegatronDistributedFusedAdam obj's underlying member method code with example usage: obj.grad_norm(parameters, norm_type, force) and returns: super