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 hasattr(module, "weight"):
              if module.weight.data.is_meta:
                  # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
                  module.weight = torch.nn.parameter.Parameter(
                      data=torch.empty_like(
                          module.weight.data, device="cpu"
                      ),
                      requires_grad=module.weight.data.requires_grad,
                  )
              if "query_key_value" in prefix:
                  module.weight = self.mp_replace.strided_copy(
                      module.weight.data,
                      state_dict[prefix + "weight"],
                      num_splits=3,
                  )
              else:
                  module.weight = self.mp_replace.copy(
                      module.weight.data, state_dict[prefix + "weight"]
                  )
          else:
              if module.norm.weight.data.is_meta:
                  # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
                  module.norm.weight = torch.nn.parameter.Parameter(
                      data=torch.empty_like(
                          module.norm.weight.data, device="cpu"
                      ),
                      requires_grad=module.norm.weight.data.requires_grad,
                  )
              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"):
                  if module.norm.bias.data.is_meta:
                      # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
                      module.norm.bias = torch.nn.parameter.Parameter(
                          data=torch.empty_like(
                              module.norm.bias.data, device="cpu"
                          ),
                          requires_grad=module.norm.bias.data.requires_grad,
                      )
                  module.norm.bias = self.mp_replace.copy(
                      module.norm.bias, state_dict[prefix + "bias"]
                  )
              else:
                  if module.bias.data.is_meta:
                      # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
                      module.bias = torch.nn.parameter.Parameter(
                          data=torch.empty_like(
                              module.bias.data, device="cpu"
                          ),
                          requires_grad=module.bias.data.requires_grad,
                      )
                  data = state_dict[prefix + "bias"]
                  data = data.to(
                      get_accelerator().current_device_name()
                  )
                  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)
  
      embedding_weight = None
  
      for n, p in r_module.named_parameters():
          if (
              "word_embeddings." in n
              or "embed_tokens." in n
              or "wte." in n
          ):
              embedding_weight = p
      if (
          embedding_weight is not None
          and hasattr(r_module, "lm_head")
          and hasattr(r_module.lm_head, "weight")
          and r_module.lm_head.weight.is_meta
      ):
          r_module.lm_head.weight = embedding_weight
  
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 import replace_transformer_layer, generic_injection
  from ..comm.comm import init_distributed
  from ..pipe import PipelineModule
  from ..moe.utils import has_moe_layers
  from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
  from deepspeed.accelerator import get_accelerator
  from ..module_inject.policy import TransformerPolicy
  from ..module_inject.auto_tp import AutoTP
  from ..module_inject.replace_policy import generic_policies
  from ..module_inject.auto_tp_model_utils import build_bloom_alibi_tensor, build_mpt_atten_bias_tensor, build_mpt_alibi_tensor, get_alibi_mask
  from ..ops.transformer.inference.ds_attention import DeepSpeedSelfAttention
  from ..model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
  from torch import nn
label_desc: Use image node_id 14 for calling the InferenceEngine obj's underlying member method code with example usage: obj.load_model_with_checkpoint(r_module) without return types