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: Load parameter state (i.e., parameter & optimizer tensors).
  
  This method performs the reverse of save_parameter_state():
  - Load world buffers from disk (i.e., distrib_opt.pt).
  - Scatter contiguous buffers from DP rank 0 to each DP rank (each DP
    rank receives its relevant subset of the world buffers).
  - For each DP rank, copy param & optimizer shards from contiguous CPU
    buffers. (e.g., one buffer each for main_param, exp_avg, and
    exp_avg_sq).
docstr_tok: ["Load","parameter","state","(","i.e.",",","parameter","&","optimizer","tensors",")",".","This","method","performs","the","reverse","of","save_parameter_state","(",")",":","-","Load","world","buffers","from","disk","(","i.e.",",","distrib_opt.pt",")",".","-","Scatter","contiguous","buffers","from","DP","rank","0","to","each","DP","rank","(","each","DP","rank","receives","its","relevant","subset","of","the","world","buffers",")",".","-","For","each","DP","rank",",","copy","param","&","optimizer","shards","from","contiguous","CPU","buffers",".","(","e.g.",",","one","buffer","each","for","main_param",",","exp_avg",",","and","exp_avg_sq",")","."]
code: def load_parameter_state(self, filename):
      """Load parameter state (i.e., parameter & optimizer tensors).
  
      This method performs the reverse of save_parameter_state():
      - Load world buffers from disk (i.e., distrib_opt.pt).
      - Scatter contiguous buffers from DP rank 0 to each DP rank (each DP
        rank receives its relevant subset of the world buffers).
      - For each DP rank, copy param & optimizer shards from contiguous CPU
        buffers. (e.g., one buffer each for main_param, exp_avg, and
        exp_avg_sq).
      """
  
      # Data parallelism variables.
      data_parallel_world_size = mpu.get_data_parallel_world_size(
          with_context_parallel=True
      )
      data_parallel_rank = mpu.get_data_parallel_rank(
          with_context_parallel=True
      )
      data_parallel_group_gloo = mpu.get_data_parallel_group_gloo(
          with_context_parallel=True
      )
      data_parallel_global_ranks = list(
          mpu._DATA_PARALLEL_GLOBAL_RANKS_WITH_CP
      )
  
      # Load on DP rank 0.
      if data_parallel_rank == 0:
          loaded_state = torch.load(filename)
          if "per_bucket_numel_unpadded" in loaded_state:
              per_bucket_numel_unpadded_in_checkpoint = loaded_state[
                  "per_bucket_numel_unpadded"
              ]
              assert (
                  self.per_bucket_numel_unpadded
                  == per_bucket_numel_unpadded_in_checkpoint
              ), (
                  f"Number of unpadded elements in each bucket need to be the same in current run "
                  f"({self.per_bucket_numel_unpadded}) and checkpoint "
                  f"({per_bucket_numel_unpadded_in_checkpoint})"
              )
  
      # Scatter tensors to all DP ranks.
      for model_idx, gbuf_range_maps in enumerate(
          self.model_gbuf_ranges
      ):
          for (
              dtype,
              gbuf_range_map_for_all_buckets,
          ) in gbuf_range_maps.items():
              for bucket_idx, gbuf_range_map in enumerate(
                  gbuf_range_map_for_all_buckets
              ):
                  # Compute local DP contiguous shard's size.
                  model = self.models[model_idx]
                  gbuf_world_numel = (
                      model.grad_buffers[dtype]
                      .buckets[bucket_idx]
                      .data.numel()
                  )
                  assert (
                      gbuf_world_numel
                      == self.per_bucket_numel[model_idx][dtype][
                          bucket_idx
                      ]
                  )
                  assert (
                      gbuf_world_numel % data_parallel_world_size == 0
                  )
                  gbuf_local_numel = (
                      gbuf_world_numel // data_parallel_world_size
                  )
  
                  # Contiguous local shards (received from DP rank 0).
                  local_shards = {
                      key: torch.empty(
                          (gbuf_local_numel,),
                          dtype=torch.float32,
                          device="cpu",
                      )
                      for key in ("param", "exp_avg", "exp_avg_sq")
                  }
  
                  # Scatter local shards from DP rank 0.
                  for key, recv_tensor in local_shards.items():
                      # Scatter tensor list.
                      if data_parallel_rank == 0:
                          world_tensor_for_all_buckets = loaded_state[
                              model_idx
                          ][dtype][key]
                          if not isinstance(
                              world_tensor_for_all_buckets, list
                          ):
                              world_tensor_for_all_buckets = [
                                  world_tensor_for_all_buckets
                              ]
                          assert bucket_idx < len(
                              world_tensor_for_all_buckets
                          ), (
                              f"Trying to load state for bucket_id {bucket_idx} (out of "
                              f"{len(gbuf_range_map_for_all_buckets)} buckets) from checkpoint; "
                              f"checkpoint only has {len(world_tensor_for_all_buckets)} bucket(s)"
                          )
                          # This tensor might be bigger or smaller than expected (depending on
                          # relative sizes of per_bucket_numel_in_checkpoint and self.per_bucket_numel).
                          world_tensor = world_tensor_for_all_buckets[
                              bucket_idx
                          ]
                          if "per_bucket_numel" in loaded_state:
                              numel_in_checkpoint = loaded_state[
                                  "per_bucket_numel"
                              ][model_idx][dtype][bucket_idx]
                              numel = self.per_bucket_numel[model_idx][
                                  dtype
                              ][bucket_idx]
                              numel_unpadded = (
                                  self.per_bucket_numel_unpadded[
                                      model_idx
                                  ][dtype][bucket_idx]
                              )
                              assert (
                                  world_tensor.numel()
                                  == numel_in_checkpoint
                              )
                              assert (
                                  numel_unpadded <= world_tensor.numel()
                              ), (
                                  "True number of elements should be fewer than number of elements in "
                                  "checkpoint tensor"
                              )
                              if world_tensor.numel() > numel:
                                  # Truncate extra values, which are padding anyway.
                                  print_rank_0(
                                      f"Truncating extra values from checkpoint (numel_in_checkpoint={numel_in_checkpoint}, "
                                      f"numel={numel}, numel_unpadded={numel_unpadded})"
                                  )
                                  world_tensor = world_tensor[:numel]
                              elif world_tensor.numel() < numel:
                                  # In this case, numel > world_tensor.numel() (which is numel_in_checkpoint).
                                  # Create new tensor with right number of values, then copy and use new tensor.
                                  print_rank_0(
                                      f"Expanding tensor from checkpoint (numel_in_checkpoint={numel_in_checkpoint}, "
                                      f"numel={numel}, numel_unpadded={numel_unpadded})"
                                  )
                                  world_tensor_reshaped = torch.empty(
                                      (numel,),
                                      dtype=world_tensor.dtype,
                                      device=world_tensor.device,
                                  )
                                  world_tensor_reshaped[
                                      :numel_in_checkpoint
                                  ].copy_(world_tensor)
                                  world_tensor = world_tensor_reshaped
                          else:
                              print_rank_0(
                                  "***WARNING*** Using older checkpoint so skipping padding checks"
                              )
                          gbuf_start_idxs = list(
                              range(
                                  0, gbuf_world_numel, gbuf_local_numel
                              )
                          )
                          send_tensors = [
                              world_tensor[i : (i + gbuf_local_numel)]
                              for i in gbuf_start_idxs
                          ]
                      else:
                          send_tensors = None
  
                      # Scatter.
                      torch.distributed.scatter(
                          recv_tensor,
                          send_tensors,
                          data_parallel_global_ranks[0],
                          data_parallel_group_gloo,
                      )
  
                  # Copy local contiguous shards to param/optim shards.
                  for model_param, param_range_map in gbuf_range_map[
                      "param_map"
                  ].items():
                      # Main param & optimizer states.
                      (
                          group_index,
                          group_order,
                      ) = self.model_param_group_index_map[model_param]
                      main_param = self.optimizer.param_groups[
                          group_index
                      ]["params"][group_order]
                      optim_state = self.optimizer.state[main_param]
  
                      tensors = {
                          "param": main_param,
                          **optim_state,
                      }
  
                      # Copy states into contiguous shard.
                      gbuf_local_start = param_range_map[
                          "gbuf_local"
                      ].start
                      gbuf_local_end = param_range_map["gbuf_local"].end
                      for key in local_shards:
                          tensors[key].data.copy_(
                              local_shards[key][
                                  gbuf_local_start:gbuf_local_end
                              ]
                          )
  
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 13 for calling the DistributedOptimizer obj's underlying member method code with example usage: obj.load_parameter_state(filename) without return types