code: def split_state_dict(
      self,
      mp_world_size,
      mp_rank,
      quantize=False,
      quantize_bits=8,
      groups=64,
      mlp_extra_grouping=True,
  ):
      # self.sanity_check(self.ckpt_list[0])
  
      sd, num_to_split, ckpt_offset = self.get_split_state_dict(
          mp_world_size, mp_rank
      )
      ds_sd = copy.deepcopy(sd)
      new_client_sd = collections.OrderedDict()
  
      client_sd = self.get_module(sd)
  
      ckpt_ver = self.get_checkpoint_version(ds_sd)
      logger.info(f"checkpoint version: {ckpt_ver}")
  
      if quantize:
          quantizer = WeightQuantization(
              mlp_extra_grouping=mlp_extra_grouping,
              mp_size=mp_world_size,
          )
  
      for key in client_sd.keys():
          value = client_sd[key]
  
          if (
              "attention.dense.weight" in key
              or "mlp.dense_4h_to_h.weight" in key
          ):
              assert value.shape[1] % num_to_split == 0
              split_size = value.shape[1] // num_to_split
              if quantize:
                  q_vals = quantizer.Quantize(
                      [value], quantize_bits, groups, key
                  )
                  value = q_vals[0]
              new_client_sd[key] = torch.split(
                  value, split_size, dim=1
              )[ckpt_offset]
          elif "attention.query_key_value" in key:
              if quantize and "attention.query_key_value.weight" in key:
                  q_vals = quantizer.Quantize(
                      [value], quantize_bits, groups, key
                  )
                  value = q_vals[0]
              new_client_sd[key] = self.split_query_key_value(
                  value, num_to_split, ckpt_offset, ckpt_ver
              )
          elif (
              "mlp.dense_h_to_4h.weight" in key
              or "word_embeddings.weight" in key
              or "mlp.dense_h_to_4h.bias" in key
              or "final_linear.weight" in key
          ):
              assert value.shape[0] % num_to_split == 0
              split_size = value.shape[0] // num_to_split
              if quantize and "mlp.dense_h_to_4h.weight" in key:
                  q_vals = quantizer.Quantize(
                      [value], quantize_bits, groups, key
                  )
                  value = q_vals[0]
              new_client_sd[key] = torch.split(
                  value, split_size, dim=0
              )[ckpt_offset]
          else:
              new_client_sd[key] = value
  
      if quantize:
          all_scales = quantizer.merge_scales_split(num_to_split)
  
      ds_sd = self.set_module(ds_sd, new_client_sd)
  
      return ds_sd, (all_scales if quantize else None)
  
imports: import torch
  import os
  import copy
  import collections
  import json
  from abc import ABC, abstractmethod
  from deepspeed.utils import logger
  from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
  from .weight_quantizer import WeightQuantization
label_desc: Use image node_id 5 for calling the MegatronSDLoader obj's underlying member method code with example usage: obj.split_state_dict(mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping) and returns: ds_sd,  unknown