code: def get_split_state_dict(self, mp_world_size, mp_rank):
      num_ckpt = len(self.ckpt_list)
      assert (
          mp_world_size % num_ckpt == 0
      ), "Invalid checkpoints and world size for sd split"
  
      num_to_split = mp_world_size // num_ckpt
      ckpt_index = mp_rank // num_to_split
      ckpt_offset = mp_rank % num_to_split
  
      logger.info(
          f"mp_rank: {mp_rank}, ckpt_list: {self.ckpt_list[ckpt_index]}, offset: {ckpt_offset}"
      )
  
      sd = self.checkpoint_engine.load(
          self.ckpt_list[ckpt_index],
          map_location=lambda storage, loc: storage,
      )
  
      return sd, num_to_split, ckpt_offset
  
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 4 for calling the SDLoaderBase obj's underlying member method code with example usage: obj.get_split_state_dict(mp_world_size, mp_rank) and returns: sd,  num_to_split,  ckpt_offset