code: def load_model_with_checkpoint(self, r_module):
      self.mp_replace = ReplaceWithTensorSlicing(
          mp_group=self.mp_group,
          mp_size=self._config.tensor_parallel.tp_size,
      )  # , out_dim=0, in_dim=1)
      error_msgs = []
  
      def load(module, state_dict, prefix):
          args = (state_dict, prefix, {}, True, [], [], error_msgs)
          if (
              len(list(module.parameters())) > 0
              and list(module.parameters())[0].numel() == 0
          ):
              with GatheredParameters(
                  list(module.parameters(recurse=False)),
                  modifier_rank=0,
              ):
                  if dist.get_rank() == 0:
                      module._load_from_state_dict(*args)
          else:
              if hasattr(module, "weight"):
                  if "query_key_value" in prefix:
                      module.weight = self.mp_replace.qkv_copy(
                          module.weight.data,
                          state_dict[prefix + "weight"],
                      )
                  else:
                      module.weight = self.mp_replace.copy(
                          module.weight.data,
                          state_dict[prefix + "weight"],
                      )
              else:
                  module.norm.weight = self.mp_replace.copy(
                      module.norm.weight.data,
                      state_dict[prefix + "weight"],
                  )
              if prefix + "bias" in self.key_list:
                  if hasattr(module, "norm"):
                      module.norm.bias = self.mp_replace.copy(
                          module.norm.bias, state_dict[prefix + "bias"]
                      )
                  else:
                      data = state_dict[prefix + "bias"]
                      data = data.to(torch.cuda.current_device())
                      module.bias = self.mp_replace.copy(
                          module.bias, data
                      )
  
      layer_policies = {
          nn.Linear: load,
          nn.Embedding: load,
          nn.LayerNorm: load,
          LinearLayer: load,
          LinearAllreduce: load,
      }
  
      def load_module_recursive(module, prefix="", level=0):
          for name, child in module.named_children():
              if child.__class__ in layer_policies:
                  checking_key = prefix + name + "."
                  if not any(
                      checking_key in item for item in self.key_list
                  ):
                      continue
                  if (
                      len(list(child.parameters())) > 0
                      and list(child.parameters())[0].numel() == 0
                  ):
                      if len(child.weight.ds_shape) == 1:
                          child = Normalize(
                              dim=child.weight.ds_shape[-1],
                              dtype=child.weight.dtype,
                              eps=child.eps,
                          )
                          setattr(module, name, child)
                  load(child, self.sd, prefix + name + ".")
              else:
                  load_module_recursive(
                      child,
                      prefix if level == 0 else prefix + name + ".",
                      level + 1,
                  )
  
      load_module_recursive(r_module)
  
imports: import torch
  import time
  import os
  from deepspeed import comm
  from deepspeed.utils.logging import log_dist
  from torch.nn.modules import Module
  from packaging import version
  from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
  from deepspeed.utils.timer import SynchronizedWallClockTimer
  from ..runtime.state_dict_factory import SDLoaderFactory
  from ..runtime.weight_quantizer import WeightQuantization
  from ..module_inject.replace_module import replace_transformer_layer, generic_injection
  from ..comm.comm import init_distributed
  from ..pipe import PipelineModule
  from ..moe.utils import has_moe_layers
  from ..runtime.zero import GatheredParameters
  from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
  from ..module_inject.replace_policy import DSPolicy
  from torch import nn
label_desc: Use image node_id 11 for calling the InferenceEngine obj's underlying member method code with example usage: obj.load_model_with_checkpoint(r_module) without return types