class_docstr: Float16 optimizer for fp16 and bf16 data types.
  
  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.
      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,
          fp16,
          bf16,
          params_dtype,
          grad_scaler,
          models,
      )
  
      # ======================
      # main parameter stuff
      # ======================
  
      # Three groups of parameters:
      #   float16_groups: original float16 parameters
      #   fp32_from_float16_groups: fp32 copy of float16 parameters
      #   fp32_from_fp32_groups: original fp32 parameters
      self.float16_groups = []
      self.fp32_from_float16_groups = []
      self.fp32_from_fp32_groups = []
  
      # For all the groups in the original optimizer:
      for param_group in 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 i, param in enumerate(param_group["params"]):
              if param.requires_grad:
                  # float16 params:
                  if param.type() in [
                      "torch.cuda.HalfTensor",
                      "torch.cuda.BFloat16Tensor",
                  ]:
                      float16_params_this_group.append(param)
                      # Create a copy
                      main_param = param.detach().clone().float()
                      # Copy tensor model parallel attributes.
                      tensor_parallel.copy_tensor_model_parallel_attributes(
                          main_param, param
                      )
                      if hasattr(param, "shared"):
                          main_param.shared = param.shared
                      # Replace the optimizer params with the new fp32 copy.
                      param_group["params"][i] = 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"][i] = param
  
                  else:
                      raise TypeError(
                          "Wrapped parameters must be one of "
                          "torch.cuda.FloatTensor,  "
                          "torch.cuda.HalfTensor, or "
                          "torch.cuda.BFloat16Tensor. "
                          "Received {}".format(param.type())
                      )
  
          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)
  
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 Float16OptimizerWithFloat16Params object from inherited base classes: MixedPrecisionOptimizer with example: obj = Float16OptimizerWithFloat16Params(optimizer, clip_grad, log_num_zeros_in_grad, check_for_nan_in_grad, params_have_main_grad, fp16, bf16, params_dtype, grad_scaler, models)