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: Create optimizer groups.
  
  Given the set of parameter shard ranges that are owned by the current
  data-parallel (DP) rank, gather the set of parameters that will be
  used (in the method below) to create the current DP's optimizer
  groups.
docstr_tok: ["Create","optimizer","groups",".","Given","the","set","of","parameter","shard","ranges","that","are","owned","by","the","current","data-parallel","(","DP",")","rank",",","gather","the","set","of","parameters","that","will","be","used","(","in","the","method","below",")","to","create","the","current","DP","'s","optimizer","groups","."]
code: def build_optimizer_group_ranges(
      cls, param_groups, model_gbuf_ranges
  ):
      """
      Create optimizer groups.
  
      Given the set of parameter shard ranges that are owned by the current
      data-parallel (DP) rank, gather the set of parameters that will be
      used (in the method below) to create the current DP's optimizer
      groups.
      """
  
      num_groups = len(param_groups)
  
      # Param group map.
      # World param group map.
      # - Store a mapping of <model_parameter:group_index> for all parameters
      #   across all DP ranks. This is necessary because it is our first
      #   cross reference between the DDP mappings and the optimizer group
      #   parameters. This mapping only for use in the next step of building
      #   the local mapping over this DP rank's parameters.
      world_param_group_map = {}
      for group_index, group in enumerate(param_groups):
          for param in group["params"]:
              assert param.requires_grad
              world_param_group_map[param] = group_index
  
      # Optimizer group ranges & param-group mapping.
      # - Build a mapping from groups to their contained parameters, and also
      #   from parameters to their containing group index and order within
      #   the group. The group index and order are particularly important for
      #   saving and loading checkpoints.
      local_param_group_map = {}
      group_ranges = [{"params": []} for _ in param_groups]
      for model_gbuf_range_map in model_gbuf_ranges:
          for (
              dtype,
              gbuf_range_map_for_all_buckets,
          ) in model_gbuf_range_map.items():
              for gbuf_range_map in gbuf_range_map_for_all_buckets:
                  for param in gbuf_range_map["param_map"]:
                      group_index = world_param_group_map[param]
                      group_range = group_ranges[group_index]
                      group_range["params"].append(param)
                      local_param_group_map[param] = (
                          group_index,
                          len(group_range["params"]) - 1,
                      )
  
      # Squeeze zero-size group ranges.
      for group_index, group_range in enumerate(group_ranges):
          group_range["orig_group"] = param_groups[group_index]
          group_range["orig_group_idx"] = param_groups[group_index]
  
      return local_param_group_map, group_ranges
  
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 5 for calling the DistributedOptimizer obj's underlying member method code with example usage: obj.build_optimizer_group_ranges(cls, param_groups, model_gbuf_ranges) and returns: local_param_group_map,  group_ranges