class_docstr: Distributed optimizer, for all data types (fp16, bf16, and fp32).
  
  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.
docstr: See top of class definition for argument descriptions.
  
  The steps in this method create the core mapping between DDP grad
  buffers, parameters, and parameter shard ranges, that is needed for
  converting between model param indexes and main parameter shard
  indexes. This method also updates the optimizer parameter groups
  with the newly created shards.
docstr_tok: ["See","top","of","class","definition","for","argument","descriptions",".","The","steps","in","this","method","create","the","core","mapping","between","DDP","grad","buffers",",","parameters",",","and","parameter","shard","ranges",",","that","is","needed","for","converting","between","model","param","indexes","and","main","parameter","shard","indexes",".","This","method","also","updates","the","optimizer","parameter","groups","with","the","newly","created","shards","."]
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,
  ):
      """
      See top of class definition for argument descriptions.
  
      The steps in this method create the core mapping between DDP grad
      buffers, parameters, and parameter shard ranges, that is needed for
      converting between model param indexes and main parameter shard
      indexes. This method also updates the optimizer parameter groups
      with the newly created shards.
      """
  
      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,
      )
  
      assert isinstance(
          optimizer, Adam
      ), "Only Adam currently supported, due to checkpointing requirements."
  
      # Model grad buffer ranges.
      self.model_gbuf_ranges = []
      self.per_bucket_numel = []
      self.per_bucket_numel_unpadded = []
      for _, model_chunk in enumerate(self.models):
          self.per_bucket_numel.append(
              {
                  dtype: [
                      bucket.data.numel()
                      for bucket in model_chunk.grad_buffers[
                          dtype
                      ].buckets
                  ]
                  for dtype in model_chunk.grad_buffers
              }
          )
          self.per_bucket_numel_unpadded.append(
              {
                  dtype: [
                      bucket.numel_unpadded
                      for bucket in model_chunk.grad_buffers[
                          dtype
                      ].buckets
                  ]
                  for dtype in model_chunk.grad_buffers
              }
          )
          self.model_gbuf_ranges.append(
              self.build_model_gbuf_range_map(model_chunk)
          )
      self.model_param_gbuf_map = self.build_model_param_gbuf_map(
          self.model_gbuf_ranges
      )
  
      # Optimizer ranges.
      (
          self.model_param_group_index_map,
          self.opt_group_ranges,
      ) = self.build_optimizer_group_ranges(
          self.optimizer.param_groups, self.model_gbuf_ranges
      )
  
      # Allocate main param shards.
      (
          self.model_float16_groups,
          self.model_fp32_groups,
          self.shard_float16_groups,
          self.shard_fp32_groups,
          self.shard_fp32_from_float16_groups,
      ) = self.build_model_and_main_param_groups(
          self.model_gbuf_ranges,
          self.model_param_gbuf_map,
          self.opt_group_ranges,
      )
  
      # Initialize param buffers.
      # - These are views on the DDP model's grad buffers, that share
      #   storage & have their own dtype. This is safe because the param
      #   dtype size is always <= grad dtype size.
      self.param_buffers = []
      for model_index, model in enumerate(self.models):
          current_param_buffers = {}
          for dtype, grad_buffer in model.grad_buffers.items():
              size_ratio = (
                  torch.finfo(dtype).bits
                  // torch.finfo(params_dtype).bits
              )
              current_param_buffers[dtype] = []
              for bucket in grad_buffer.buckets:
                  # Handle older/newer method for getting untyped storage.
                  try:
                      storage = bucket.data.untyped_storage()
                  except:
                      try:
                          storage = bucket.data.storage()._untyped()
                      except:
                          storage = bucket.data.storage().untyped()
  
                  # Typed param buffer.
                  param_buffer = torch.tensor(
                      storage,
                      dtype=params_dtype,
                      device=bucket.data.device,
                  )
  
                  # .storage() ignores views / slices, so param_buffer now points to the start
                  # of the grad_buffer instead of to the start of each bucket. As a result,
                  # add bucket.offset to make sure param_buffers point to the right region of
                  # memory.
                  # Since we want the start of each bucket's param_buffer to coincide with the
                  # start of the same bucket's grad_buffer (this ensures that zeroing the grad
                  # buffer does not zero out params in the param_buffer before they are copied
                  # into the model_params), multiply the offset by the size ratio of grads and
                  # params.
                  offset = bucket.offset * size_ratio
                  param_buffer = param_buffer[
                      offset : offset + bucket.data.numel()
                  ]
                  assert (
                      param_buffer.data_ptr() == bucket.data.data_ptr()
                  ), "param_buffer and grad_buffer for same bucket should start at the same byte address"
                  assert (
                      param_buffer.numel() == bucket.data.numel()
                  ), "param_buffer and grad_buffer for same bucket should have the same number of elements"
                  current_param_buffers[dtype].append(param_buffer)
          self.param_buffers.append(current_param_buffers)
  
      # Now construct data structures to manage all-gather handles.
      self.all_gather_handles = []
      self.all_gather_handle_index_to_bucket_index_map = []
      self.model_index_to_all_gather_handle_index_map = {}
      self.param_to_all_gather_handle_index_map = {}
      self.param_buffer_copied = []
  
      self.pbuf_view_items = self.get_model_param_buffer_dp_views()
      for (
          model_index,
          dtype,
          bucket_index,
          _,
          _,
      ) in self.pbuf_view_items:
          self.all_gather_handle_index_to_bucket_index_map.append(
              (model_index, dtype, bucket_index)
          )
          all_gather_handle_index = (
              len(self.all_gather_handle_index_to_bucket_index_map) - 1
          )
  
          # Store all all_gather_handle_indices relevant to a particular model chunk.
          if (
              model_index
              not in self.model_index_to_all_gather_handle_index_map
          ):
              self.model_index_to_all_gather_handle_index_map[
                  model_index
              ] = []
          self.model_index_to_all_gather_handle_index_map[
              model_index
          ].append(all_gather_handle_index)
  
          for param in (
              self.models[model_index]
              .grad_buffers[dtype]
              .buckets[bucket_index]
              .params_list
          ):
              self.param_to_all_gather_handle_index_map[
                  param
              ] = all_gather_handle_index
          self.param_buffer_copied.append(False)
      self.num_all_gather_handles = len(
          self.all_gather_handle_index_to_bucket_index_map
      )
  
      self.overlap_param_gather = get_args().overlap_param_gather
      if self.overlap_param_gather:
          self.remove_pre_hook_handle = (
              torch.nn.modules.module.register_module_forward_pre_hook(
                  self._make_forward_pre_hook()
              )
          )
      else:
          self.remove_pre_hook_handle = None
  
      self.update_successful = False
  
      # Update optimizer groups.
      # - Also, leverage state_dict() and load_state_dict() to
      #   recast preexisting per-param state tensors.
      self.optimizer.param_groups = [
          g["orig_group"] for g in self.opt_group_ranges
      ]
      self.optimizer.load_state_dict(self.optimizer.state_dict())
  
imports: from apex.optimizers import FusedAdam
  import math
  import torch
  from megatron import get_args
  from megatron import get_timers
  from megatron import print_rank_0
  from megatron.core import mpu, tensor_parallel
  from .optimizer import MixedPrecisionOptimizer, _zero_grad_group_helper
  from .utils import shard_buffer
label_desc: Use image node_id 7 to create a new DistributedOptimizer object from inherited base classes: MixedPrecisionOptimizer with example: obj = DistributedOptimizer(optimizer, clip_grad, log_num_zeros_in_grad, check_for_nan_in_grad, params_have_main_grad, fp16, bf16, params_dtype, grad_scaler, models)