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: All-gather updated model params.
  
  The DDP's param buffer is used for the all-gather, and thus no
  tensors are dynamically allocated. After the all-gather, the params
  can be copied from the param buffer to the param.
docstr_tok: ["All-gather","updated","model","params",".","The","DDP","'s","param","buffer","is","used","for","the","all-gather",",","and","thus","no","tensors","are","dynamically","allocated",".","After","the","all-gather",",","the","params","can","be","copied","from","the","param","buffer","to","the","param","."]
code: def _dispatch_gather_model_params(self, all_gather_handle_index):
      """
      All-gather updated model params.
  
      The DDP's param buffer is used for the all-gather, and thus no
      tensors are dynamically allocated. After the all-gather, the params
      can be copied from the param buffer to the param.
      """
      if self.update_successful:
          data_parallel_rank = mpu.get_data_parallel_rank(
              with_context_parallel=True
          )
          data_parallel_group = mpu.get_data_parallel_group(
              with_context_parallel=True
          )
  
          # All-gather updated main params.
          # All param_buf views are guaranteed to have the same number of elements
          # across all data-parallel ranks, due to padding (done in grad_buffer.py),
          # and extended to the param_bufs. Thus, all sub-views will have consistent
          # start / end indexes across data-parallel ranks.
          (
              model_index,
              dtype,
              bucket_index,
              pbuf,
              pbuf_views,
          ) = self.pbuf_view_items[all_gather_handle_index]
          assert all_gather_handle_index == len(self.all_gather_handles)
          all_gather_handle = torch.distributed._all_gather_base(
              pbuf,
              pbuf_views[data_parallel_rank],
              group=data_parallel_group,
              async_op=self.overlap_param_gather,
          )
          self.all_gather_handles.append(all_gather_handle)
          assert self.all_gather_handle_index_to_bucket_index_map[
              all_gather_handle_index
          ] == (model_index, dtype, bucket_index)
          self.param_buffer_copied.append(False)
  
      if not self.overlap_param_gather:
          self._copy_params_from_param_buffer(all_gather_handle_index)
  
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 16 for calling the DistributedOptimizer obj's underlying member method code with example usage: obj._dispatch_gather_model_params(all_gather_handle_index) without return types