class_docstr: Wrapper class that supports NeMo-Megatron optimizations
  
  Performs gradient clipping, unscaling, and optimizer step.
code: def step(self, closure=None, grad_scaler=None):
      # Code path below assumes capturable=True and master_weights=True
      if not (self.capturable and self.master_weights):
          return super().step(closure=closure, grad_scaler=grad_scaler)
  
      loss = None
      if closure is not None:
          loss = closure()
  
      for group, group_master in zip(
          self.param_groups, self.param_groups_master
      ):
          if len(group["params"]) == 0:
              continue
          device = group["params"][0].device
          bias_correction = 1 if group["bias_correction"] else 0
          beta1, beta2 = group["betas"]
  
          # Assume same step per parameter group for simplicity
          if "step" in group:
              group["step"] += (
                  1
                  if not self.capturable
                  else (self._dummy_overflow_buf != 1).to(torch.int)
              )
          else:
              group["step"] = (
                  1
                  if not self.capturable
                  else torch.tensor([1], dtype=torch.int, device=device)
              )
  
          # Check for overflow in gradients
          found_inf = (
              grad_scaler._check_inf_per_device(self)[device]
              if grad_scaler is not None
              else torch.zeros((1,), device=device)
          )
          self._dummy_overflow_buf.copy_(found_inf)
  
          # Get gradient scaling/unscaling factors
          scale, inv_scale = None, None
          if grad_scaler:
              scale = grad_scaler._get_scale_async()
              inv_scale = scale.double().reciprocal().float()
          else:
              scale = torch.ones((1,), device=device)
              inv_scale = torch.ones((1,), device=device)
          combined_scale = inv_scale
  
          # Gradient clipping
          if self.max_norm > 0:
              # Unscale gradients and find L2 norm
              fp32_grads_for_norm = []
              fp16_grads_for_norm = []
              for p in group["params"]:
                  if p.grad is None:
                      continue
                  assert p.dtype in [
                      torch.float32,
                      torch.float16,
                  ], "Only FP32/FP16 model parameters are supported"
  
                  is_not_shared = param_is_not_shared(p)
                  is_not_tp_duplicate = (
                      param_is_not_tensor_parallel_duplicate(p)
                  )
                  if is_not_shared and is_not_tp_duplicate:
                      if p.dtype == torch.float32:
                          fp32_grads_for_norm.append(p.grad.detach())
                      else:
                          fp16_grads_for_norm.append(p.grad.detach())
  
              if fp32_grads_for_norm:
                  fp32_grad_norm, _ = multi_tensor_applier(
                      amp_C.multi_tensor_unscale_l2norm,
                      self._dummy_overflow_buf,
                      [fp32_grads_for_norm],
                      inv_scale,
                      False,
                  )
              else:
                  fp32_grad_norm = torch.tensor(
                      [0.0], dtype=torch.float32, device=device
                  )
  
              if fp16_grads_for_norm:
                  fp16_grad_norm, _ = multi_tensor_applier(
                      amp_C.multi_tensor_unscale_l2norm,
                      self._dummy_overflow_buf,
                      [fp16_grads_for_norm],
                      inv_scale,
                      False,
                  )
              else:
                  fp16_grad_norm = torch.tensor(
                      [0.0], dtype=torch.float32, device=device
                  )
  
              # Prep L2 norm for allreduce
              total_norm = (
                  fp32_grad_norm**self.norm_type
                  + fp16_grad_norm**self.norm_type
              ).squeeze()
  
              # Allreduce L2 norm across model-parallel GPUs
              torch.distributed.all_reduce(
                  total_norm,
                  op=torch.distributed.ReduceOp.SUM,
                  group=parallel_state.get_model_parallel_group(),
              )
              total_norm = total_norm ** (1.0 / self.norm_type)
  
              # Combine unscaling factor with clip coefficient
              clip_coeff = self.max_norm / (total_norm + 1.0e-6)
              clip_coeff_clamped = torch.clamp(clip_coeff, max=1.0)
              combined_scale = (
                  clip_coeff_clamped * combined_scale
              )  # Potential issue with associativity?
  
          # Create lists for multi-tensor apply
          g_16, p_16, m_16, v_16 = [], [], [], []
          g_32, p_32, m_32, v_32 = [], [], [], []
          p_16_master = []
          p_32_master = []
  
          for p, p_master in zip(
              group["params"], group_master["params"]
          ):
              if p.grad is None:
                  continue
              if p.grad.data.is_sparse:
                  raise RuntimeError(
                      "MegatronFusedAdam does not support sparse gradients, please consider SparseAdam instead"
                  )
  
              state = self.state[p]
              # State initialization
              if len(state) == 0:
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(p.data).float()
                  # Exponential moving average of squared gradient values
                  state["exp_avg_sq"] = torch.zeros_like(p.data).float()
  
              if p.dtype == torch.float16:
                  p_16_master.append(p_master.data)
                  g_16.append(p.grad.data)
                  p_16.append(p.data)
                  m_16.append(state["exp_avg"])
                  v_16.append(state["exp_avg_sq"])
              elif p.dtype == torch.float32:
                  p_32_master.append(p_master.data)
                  g_32.append(p.grad.data)
                  p_32.append(p.data)
                  m_32.append(state["exp_avg"])
                  v_32.append(state["exp_avg_sq"])
              else:
                  raise RuntimeError(
                      "MegatronFusedAdam only supports fp16 and fp32."
                  )
  
          if len(g_16) > 0:
              multi_tensor_applier(
                  self.multi_tensor_adam_capturable_master,
                  self._dummy_overflow_buf,
                  [g_16, p_16, m_16, v_16, p_16_master],
                  group["lr"],
                  beta1,
                  beta2,
                  group["eps"],
                  group["step"],
                  self.adam_w_mode,
                  bias_correction,
                  group["weight_decay"],
                  combined_scale,
              )
  
          if len(g_32) > 0:
              multi_tensor_applier(
                  self.multi_tensor_adam_capturable_master,
                  self._dummy_overflow_buf,
                  [g_32, p_32, m_32, v_32, p_32_master],
                  group["lr"],
                  beta1,
                  beta2,
                  group["eps"],
                  group["step"],
                  self.adam_w_mode,
                  bias_correction,
                  group["weight_decay"],
                  combined_scale,
              )
  
      return loss
  
imports: import amp_C
  import torch
  from nemo.collections.nlp.modules.common.megatron.module import param_is_not_shared
label_desc: Use image node_id 2 for calling the MegatronFusedAdam obj's underlying member method code with example usage: obj.step(closure, grad_scaler) and returns: loss, super