class_docstr: Base class for both the float-16 and the distributed optimizer.
  
  Arguments:
      optimizer: base optimizer such as Adam or SGD
      clip_grad: clip gradeints with this global L2 norm. Note
          that clipping is ignored if clip_grad == 0
      log_num_zeros_in_grad: return number of zeros in the gradients.
      check_for_nan_in_grad: check if gradients have a NaN.
      params_have_main_grad: flag indicating if parameters have
          a `main_grad` field. If this is set, we are assuming
          that the model parameters are store in the `main_grad`
          field instead of the typical `grad` field. This happens
          for the DDP cases where there is a continuous buffer
          holding the gradients. For example for bfloat16, we want
          to do gradient accumulation and all-reduces in float32
          and as a result we store those gradients in the main_grad.
          Note that main grad is not necessarily in float32.
      fp16: if true, the model is running in fp16.
      bf16: if true, the model is running in bfloat16.
      params_dtype: used by distributed optimizer.
      grad_scaler: used for scaling gradients. Note that this can be
          None. This case happens when `bf16 = True` and we don't
          use any loss scale. Note that for `bf16 = True`, we can have
          a constnat gradient scaler. Also for `bf16 = False`, we
          always require a grad scaler.
      models: list of models (i.e., the virtual pipelining models). This
          is used by the distributed optimizer for mapping parameters.
code: def __init__(
      self,
      optimizer,
      clip_grad,
      log_num_zeros_in_grad,
      check_for_nan_in_grad,
      params_have_main_grad,
      fp16,
      bf16,
      params_dtype,
      grad_scaler,
      models,
  ):
      super().__init__(
          optimizer,
          clip_grad,
          log_num_zeros_in_grad,
          check_for_nan_in_grad,
          params_have_main_grad,
          models,
      )
  
      self.fp16 = fp16
      self.bf16 = bf16
      self.params_dtype = params_dtype
      self.grad_scaler = grad_scaler
  
      # None grad scaler is only supported for bf16.
      if self.grad_scaler is None:
          assert not self.fp16, "fp16 expects a grad scaler."
  
      # Tensor used to determine if a nan/if has happend.
      # Any non-zero value indicates inf/nan.
      # Note that we keep this for the cases that grad scaler is none.
      # We still record nan/inf if we have a bfloat16 with a grad scaler.
      if self.grad_scaler:
          self.found_inf = torch.tensor(
              [0.0], dtype=torch.float, device="cuda"
          )
  
      # Dummy tensor needed for apex multi-apply tensor.
      # For bfloat, we don't have multi-tensor apply and for now
      # we set it to none so the multi-tensor apply gets ignored.
      if bf16:
          self._dummy_overflow_buf = None
      else:
          self._dummy_overflow_buf = torch.tensor(
              [0], dtype=torch.int, device="cuda"
          )
  
      # In case grad scaler is not passed, define the unity scale.
      if self.grad_scaler is None:
          self._scale_one = torch.tensor(
              [1.0], dtype=torch.float, device="cuda"
          )
  
imports: from abc import ABC
  from abc import abstractmethod
  from apex.multi_tensor_apply import multi_tensor_applier
  import amp_C
  import torch
  from megatron import get_timers
  from megatron import print_rank_0
  from megatron.core import mpu, tensor_parallel
  from megatron.model import Float16Module
  from megatron.model.module import param_is_not_shared
  from .clip_grads import clip_grad_norm_fp32, count_zeros_fp32
label_desc: Use image node_id 1 to create a new MixedPrecisionOptimizer object from inherited base classes: MegatronOptimizer with example: obj = MixedPrecisionOptimizer(optimizer, clip_grad, log_num_zeros_in_grad, check_for_nan_in_grad, params_have_main_grad, fp16, bf16, params_dtype, grad_scaler, models)