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()
        data_parallel_rank = mpu.get_data_parallel_rank()
        data_parallel_group_gloo = mpu.get_data_parallel_group_gloo()
        data_parallel_global_ranks = list(mpu._DATA_PARALLEL_GLOBAL_RANKS)

        # Load on DP rank 0.
        if data_parallel_rank == 0:
            loaded_state = torch.load(filename)

        # Scatter tensors to all DP ranks.
        for model_idx, gbuf_range_maps in enumerate(self.model_gbuf_ranges):
            for dtype, gbuf_range_map in gbuf_range_maps.items():

                # Compute local DP contiguous shard's size.
                model = self.models[model_idx]
                gbuf_world_numel = model._grad_buffers[dtype].numel_padded
                gbuf_local_numel = int(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 = loaded_state[model_idx][dtype][key]
                        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])
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).
        
training/DeepSpeed-Domino/domino/optimizer/distrib_optimizer.py