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: Build mapping between params and their grad buffers.
  
  This method does the initial setup for the method above. This setup
  includes determining the shard ranges into the DDP's grad buffer for
  each data-parallel (DP) rank. Each DP rank keeps range info for
  all other DP ranks, for the purpose of creating args for
  reduce-scatter and all-gather.
docstr_tok: ["Build","mapping","between","params","and","their","grad","buffers",".","This","method","does","the","initial","setup","for","the","method","above",".","This","setup","includes","determining","the","shard","ranges","into","the","DDP","'s","grad","buffer","for","each","data-parallel","(","DP",")","rank",".","Each","DP","rank","keeps","range","info","for","all","other","DP","ranks",",","for","the","purpose","of","creating","args","for","reduce-scatter","and","all-gather","."]
code: def build_model_gbuf_range(cls, model, dtype, bucket_index):
      """
      Build mapping between params and their grad buffers.
  
      This method does the initial setup for the method above. This setup
      includes determining the shard ranges into the DDP's grad buffer for
      each data-parallel (DP) rank. Each DP rank keeps range info for
      all other DP ranks, for the purpose of creating args for
      reduce-scatter and all-gather.
      """
  
      data_parallel_rank = mpu.get_data_parallel_rank(
          with_context_parallel=True
      )
      data_parallel_world_size = mpu.get_data_parallel_world_size(
          with_context_parallel=True
      )
  
      bucket = model.grad_buffers[dtype].buckets[bucket_index]
      bucket_buffer = bucket.data
      gbuf_size = bucket_buffer.numel()
      assert (
          gbuf_size % data_parallel_world_size == 0
      ), f"Each bucket's buffer size should be divisible by {data_parallel_world_size}"
      max_gbuf_range_size = gbuf_size // data_parallel_world_size
  
      # All world ranges (i.e., across all data parallel ranks).
      gbuf_world_all_ranges = []
      for r in range(data_parallel_world_size):
          # Compute start of chunk in this bucket.
          gbuf_world_start = r * max_gbuf_range_size
          gbuf_world_end = min(
              gbuf_size, gbuf_world_start + max_gbuf_range_size
          )
          # Add bucket's offset in grad buffer.
          gbuf_world_range = Range(
              gbuf_world_start + bucket.offset,
              gbuf_world_end + bucket.offset,
          )
          gbuf_world_all_ranges.append(gbuf_world_range)
  
      # Local DP's ranges.
      gbuf_world_range = gbuf_world_all_ranges[data_parallel_rank]
  
      # Get each param's ranges.
      param_range_map = cls.build_model_gbuf_param_range_map(
          model, dtype, gbuf_world_range, bucket.offset
      )
  
      # Group into dict.
      data = {
          "param_map": param_range_map,
      }
  
      return data
  
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 2 for calling the DistributedOptimizer obj's underlying member method code with example usage: obj.build_model_gbuf_range(cls, model, dtype, bucket_index) and returns: data