class_docstr: Float16 optimizer wrapper for half precision (fp16 and bf16) data types.
  This optimizer wrapper holds main parameters and gradients in fp32 to support
  stable convergence.
  
  Arguments:
      optimizer: base optimizer such as Adam or SGD.
      fp32_grad_accum: to enable the use of fp32 in gradient accumulation and allreduce.
      contiguous_grad_bucket: to enable allocating the master gradients in the 
          contiguous memory space to reduce memory fragmentation.
      async_grad_allreduce: enable asynchronous gradient allreduce that is executed
          along with the training step backprop.
code: def __init__(
      self,
      optimizer,
      fp32_grad_accum=False,
      contiguous_grad_bucket=False,
      async_grad_allreduce=False,
      grad_div_ar_fusion=True,
      grad_allreduce_chunk_size_mb=0,
  ):
      if not HAVE_APEX:
          raise ImportError(
              "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
          )
  
      if not HAVE_MEGATRON_CORE:
          raise ImportError(
              "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
          )
  
      self.optimizer = optimizer
      assert self.optimizer, "no optimizer is provided."
      if contiguous_grad_bucket:
          assert (
              fp32_grad_accum
          ), "contiguous gradient buffer assumes using fp32 grad."
      if async_grad_allreduce:
          assert fp32_grad_accum, (
              "async allreduce applies to master gradients only, "
              "which is supposed to be accumulated after grad op."
          )
          assert contiguous_grad_bucket, (
              "currently async_grad_allreduce is supported only "
              "with contiguous_grad_bucket."
          )
  
      self._fp32_grad_accum = fp32_grad_accum
      self._contiguous_grad_bucket = contiguous_grad_bucket
  
      # used with tensor parallel only (no pipeline parallelism)
      # be careful, weight update cannot start until all async grad AR works are done
      self._async_grad_allreduce = (
          async_grad_allreduce
          and get_data_parallel_world_size(with_context_parallel=True)
          > 1
      )
      self._grad_divisor = 1 / get_data_parallel_world_size(
          with_context_parallel=True
      )
  
      if self._async_grad_allreduce:
          # use @no_sync to disable backward grad sync during gradient accumulation
          self._require_backward_grad_sync = True
          self._grad_div_ar_fusion = grad_div_ar_fusion
          self._grad_allreduce_chunk_size_mb = (
              grad_allreduce_chunk_size_mb
          )
      else:
          self._require_backward_grad_sync = False
          self._grad_div_ar_fusion = False
          self._grad_allreduce_chunk_size_mb = 0
  
      # Dummy tensor needed for apex multi-apply tensor.
      self._dummy_overflow_buf = None
  
      # Create persistent buffers for main gradients in contiguous memory space
      # - Chunked element-wise and allreduce ops without creating a temporary buffer for merged operation
      # - Low memory fragmentation
      self._main_grad_buffers = None
      if self._contiguous_grad_bucket:
          self._main_grad_buffers = {}
          # get the size of buffers
          num_elements = {}
          for i, param_group in enumerate(self.optimizer.param_groups):
              for param in param_group["params"]:
                  if param.requires_grad:
                      num_elements[i] = (
                          num_elements.get(i, 0) + param.data.nelement()
                      )
  
              # Allocate gradient memory buffers for each data type
              if any(
                  param.requires_grad for param in param_group["params"]
              ):
                  self._main_grad_buffers[i] = GradBucket(
                      num_elements[i],
                      self._grad_allreduce_chunk_size_mb,
                  )
  
      # Three groups of parameters:
      self.float16_groups = []  # original float16 parameters
      self.fp32_from_float16_groups = (
          []
      )  # fp32 copy of float16 parameters
      self.fp32_from_fp32_groups = []  # original fp32 parameters
  
      # gradient function hooks
      if self._fp32_grad_accum:
          self.grad_accs = []
  
      # For all the groups in the original optimizer:
      for i, param_group in enumerate(self.optimizer.param_groups):
          float16_params_this_group = []
          fp32_params_this_group = []
          fp32_from_float16_params_this_group = []
          # For all the parameters in this group:
          for j, param in enumerate(param_group["params"]):
              main_param = None
              if param.requires_grad:
                  # float16 params:
                  if param.type() in [
                      "torch.cuda.HalfTensor",
                      "torch.cuda.BFloat16Tensor",
                  ]:
                      float16_params_this_group.append(param)
  
                      # Allocate the main parameter
                      main_param = param.detach().clone().float()
  
                      # Copy tensor model parallel attributes.
                      copy_tensor_model_parallel_attributes(
                          main_param, param
                      )
                      if hasattr(param, "shared"):
                          main_param.shared = param.shared
  
                      # Assign the grad buffer offset to main parameters
                      if self._contiguous_grad_bucket:
                          num_elements[i] -= param.data.nelement()
                          (
                              main_param.grad,
                              grad_chunk_info,
                          ) = self._main_grad_buffers[i].get(
                              param.data.shape, num_elements[i]
                          )
                          # Add a pointer to main_grad in model param for first-last stage embedding param reduction
                          param.main_grad = main_param.grad
  
                      # Replace the optimizer params with the new fp32 copy.
                      param_group["params"][j] = main_param
                      fp32_from_float16_params_this_group.append(
                          main_param
                      )
                      # Reset existing state dict key to the new main param.
                      if param in self.optimizer.state:
                          self.optimizer.state[
                              main_param
                          ] = self.optimizer.state.pop(param)
                  # fp32 params.
                  elif param.type() == "torch.cuda.FloatTensor":
                      fp32_params_this_group.append(param)
                      param_group["params"][j] = param
  
                  else:
                      raise TypeError(
                          "Wrapped parameters must be one of "
                          "torch.cuda.FloatTensor,  "
                          "torch.cuda.HalfTensor, or "
                          "torch.cuda.BFloat16Tensor. "
                          "Received {}".format(param.type())
                      )
  
              # Add gradient accumulation hook for fp32 grad accumulation
              if (
                  main_param is not None
                  and self._fp32_grad_accum
                  and param.requires_grad
              ):
                  # Expand so we get access to grad_fn
                  param_tmp = param.expand_as(param)
                  # Get the gradient accumulator function.
                  grad_acc = param_tmp.grad_fn.next_functions[0][0]
                  grad_acc.register_hook(
                      self._make_param_hook(
                          param, main_param, i, grad_chunk_info
                      )
                  )
                  self.grad_accs.append(grad_acc)
  
          self.float16_groups.append(float16_params_this_group)
          self.fp32_from_float16_groups.append(
              fp32_from_float16_params_this_group
          )
          self.fp32_from_fp32_groups.append(fp32_params_this_group)
  
      # init exp_avg and exp_avg_sq before loading optimizer state, needed for dist checkpointing
      self._init_opt_state()
  
      # Leverage state_dict() and load_state_dict() to
      # recast preexisting per-param state tensors
      self.optimizer.load_state_dict(self.optimizer.state_dict())
  
imports: from contextlib import contextmanager
  import torch
  from nemo.utils import logging
label_desc: Use image node_id 1 to create a new MainParamsOptimizerWrapper object from inherited base classes: torch.optim with example: obj = MainParamsOptimizerWrapper(optimizer, fp32_grad_accum, contiguous_grad_bucket, async_grad_allreduce, grad_div_ar_fusion, grad_allreduce_chunk_size_mb)