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: Copy params from param_buffer to model_params.
docstr_tok: ["Copy","params","from","param_buffer","to","model_params","."]
code: def _copy_params_from_param_buffer(self, all_gather_handle_index):
      """
      Copy params from param_buffer to model_params.
      """
      (
          model_index,
          dtype,
          bucket_index,
      ) = self.all_gather_handle_index_to_bucket_index_map[
          all_gather_handle_index
      ]
      model = self.models[model_index]
      if self.update_successful:
          # Copy from param buffer to each param.
          param_map = model.grad_buffer_param_index_map[dtype]
          for param, (
              buf_start,
              buf_end,
              bucket_index_in_param_map,
          ) in param_map.items():
              if bucket_index == bucket_index_in_param_map:
                  bucket_offset = (
                      model.grad_buffers[dtype]
                      .buckets[bucket_index]
                      .offset
                  )
                  param_buf = self.param_buffers[model_index][dtype][
                      bucket_index
                  ]
                  # buf_start and buf_end store position of this parameter in the full grad_buffer,
                  # so need to adjust these indices (by subtracting out bucket_offset) since we
                  # have independent param_bufs for each bucket.
                  param_buf_shard = param_buf[
                      buf_start
                      - bucket_offset : buf_end
                      - bucket_offset
                  ]
                  assert (
                      param.data.nelement()
                      == param_buf_shard.nelement()
                  )
                  param.view(-1).detach().copy_(param_buf_shard)
  
      # Zero out the grad buffer in preparation for next set of fwd / bwd passes after copy
      # completes (since param_buffer and grad_buffer are shared for each bucket).
      param_buf = self.param_buffers[model_index][dtype][bucket_index]
      grad_buf = model.grad_buffers[dtype].buckets[bucket_index].data
      assert param_buf.data_ptr() == grad_buf.data_ptr()
      grad_buf.zero_()
  
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 20 for calling the DistributedOptimizer obj's underlying member method code with example usage: obj._copy_params_from_param_buffer(all_gather_handle_index) without return types