code: def load_model_with_checkpoint(
      r_module,
      sd,
      mp_replace,
      ckpt_type,
      weight_quantizer=None,
      rank=0,
      param_names=None,
      transformer_config=None,
      megatron_v2=False,
  ):
      error_msgs = []
  
      def transpose(data):
          with torch.no_grad():
              data = data.contiguous()
              data1 = data.transpose(-1, -2).reshape(-1)
              data.reshape(-1).copy_(data1)
              data1 = None
          return data.reshape(data.shape[-1], data.shape[-2])
  
      def load(module, prefix):
          args = (sd[0], 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,
              ):
                  module._load_from_sd(*args)
          else:
              if hasattr(module, "weight"):
                  module.weight = mp_replace.copy(
                      module.weight.data, sd[0][prefix + "weight"]
                  )
              if prefix + "bias" in sd[0].keys():
                  module.bias = mp_replace.copy(
                      module.bias.data, sd[0][prefix + "bias"]
                  )
          args = None
          gc.collect()
  
      def load_transformer_layer(module, prefix):
          if ckpt_type == "tp":
  
              def load_parameters(module, prefix):
                  for n, p in module.named_parameters():
                      if prefix + n in sd[0] and len(n.split(".")) == 1:
                          if type(sd[0][prefix + n]) is list:
                              tmp_data, scale = sd[0][prefix + n]
                              tmp_data = tmp_data
                              scale = scale.to(
                                  torch.cuda.current_device()
                              )
                          else:
                              tmp_data = sd[0][prefix + n].to(
                                  torch.cuda.current_device()
                              )
                              scale = None
                          src_shape = tmp_data.shape
                          dst_shape = p.shape
                          inner_dim = (
                              1 if tmp_data.dtype == torch.int8 else 0
                          )
                          outer_dim = (
                              0 if tmp_data.dtype == torch.int8 else 1
                          )
                          if (
                              len(src_shape) == 2
                              and len(dst_shape) == 2
                          ):
                              if (
                                  src_shape[inner_dim] == dst_shape[0]
                                  and src_shape[outer_dim]
                                  == dst_shape[1]
                              ):
                                  if tmp_data.dtype != torch.int8:
                                      p = weight_quantizer.quantize(
                                          transpose(tmp_data)
                                          if weight_quantizer.q_int8
                                          else tmp_data
                                      )
                                  else:
                                      p = torch.nn.parameter.Parameter(
                                          tmp_data, requires_grad=False
                                      )
                                      p.scale = scale
                                  setattr(module, n, p)
                              else:
                                  dim = (
                                      inner_dim
                                      if src_shape[inner_dim]
                                      != dst_shape[0]
                                      else outer_dim
                                  )
                                  dim1 = (
                                      0
                                      if src_shape[inner_dim]
                                      != dst_shape[0]
                                      else 1
                                  )
                                  if src_shape[dim] > dst_shape[dim1]:
                                      weight_partition = torch.split(
                                          tmp_data,
                                          dst_shape[dim1],
                                          dim=dim,
                                      )[rank].to(
                                          torch.cuda.current_device()
                                      )
                                      assert tmp_data.dtype != torch.int8 or scale.numel() > weight_quantizer.num_groups * (
                                          rank + 1
                                      ), """ERROR: We require the quantization scales for larger TP-size when loading INT8 checkpoint!\
                                             Please use the FP16 checkpoint to generate INT8 checkpoint with the sharding parameters!"""
                                      scale = (
                                          scale.view(-1)[
                                              weight_quantizer.num_groups
                                              * (rank + 1) :
                                          ]
                                          .reshape(
                                              weight_quantizer.num_groups,
                                              -1,
                                          )
                                          .contiguous()
                                      )
                                  else:
                                      assert (
                                          tmp_data.dtype != torch.int8
                                      ), """Merging of the checkpoints are not supported when using INT8 checkpoint! \
                                            Please use a as many GPUs as TP-size for the checkpoint"""
                                      all_data = [
                                          sd[j][prefix + n]
                                          if type(sd[j][prefix + n])
                                          is list
                                          else sd[j][prefix + n].to(
                                              torch.cuda.current_device()
                                          )
                                          for j in range(len(sd))
                                      ]
                                      weight_partition = torch.cat(
                                          [
                                              ad[0].to(
                                                  torch.cuda.current_device()
                                              )
                                              if type(ad) is list
                                              else ad
                                              for ad in all_data
                                          ],
                                          dim=dim,
                                      )
                                      if tmp_data.dtype == torch.int8:
                                          scale = torch.cat(
                                              [
                                                  ad[1].to(
                                                      torch.cuda.current_device()
                                                  )
                                                  for ad in all_data
                                              ],
                                              dim=dim,
                                          )
  
                                  if tmp_data.dtype != torch.int8:
                                      weight_partition = (
                                          weight_quantizer.quantize(
                                              transpose(
                                                  weight_partition
                                              ),
                                              parallel_dim=(
                                                  0 if dim == 1 else 1
                                              ),
                                          )
                                          if weight_quantizer.q_int8
                                          else weight_quantizer.quantize(
                                              weight_partition
                                          )
                                      )
                                  else:
                                      weight_partition = (
                                          torch.nn.parameter.Parameter(
                                              weight_partition,
                                              requires_grad=False,
                                          )
                                      )
                                      weight_partition.scale = scale
                                  setattr(module, n, weight_partition)
                          else:
                              if src_shape[0] == dst_shape[0]:
                                  p.data.copy_(tmp_data)
                              else:
                                  if src_shape[0] > dst_shape[0]:
                                      bias_split = (
                                          torch.split(
                                              tmp_data, dst_shape[-1]
                                          )[rank]
                                          .to(
                                              torch.cuda.current_device()
                                          )
                                          .contiguous()
                                      )
                                      p.data.copy_(bias_split)
                                  else:
                                      p.data.copy_(
                                          torch.cat(
                                              [
                                                  sd[j][prefix + n]
                                                  for j in range(
                                                      len(sd)
                                                  )
                                              ],
                                              dim=0,
                                          )
                                          .to(
                                              torch.cuda.current_device()
                                          )
                                          .contiguous()
                                      )
  
              load_parameters(module, prefix)
              for n, child in module.named_children():
                  load_parameters(child, prefix + n + ".")
          else:
  
              def _transpose(x):
                  heads = transformer_config.heads // mp_replace.mp_size
                  attention_head_size = x.shape[-1] // heads
                  new_x_shape = x.size()[:-1] + (
                      heads,
                      attention_head_size,
                  )
                  x_1 = x.view(*new_x_shape)
                  (q, k, v) = torch.split(
                      x_1, (x_1.shape[-1] // 3), dim=(x_1.dim() - 1)
                  )
                  if len(q.shape) > 2:
                      return torch.cat(
                          (
                              q.reshape(q.shape[0], -1),
                              k.reshape(q.shape[0], -1),
                              v.reshape(q.shape[0], -1),
                          ),
                          dim=-1,
                      ).reshape(x.shape)
                  else:
                      return torch.cat(
                          (q.reshape(-1), k.reshape(-1), v.reshape(-1)),
                          dim=-1,
                      ).reshape(x.shape)
  
              # This checks if the parameter exits in the checkpoint file and maybe copies it into the corresponding destination tensor.
              # Note that not all parameters are saved in one checkpoint, that's why we always need to check if they exist!
              def maybe_copy(
                  module,
                  dst_name,
                  src_name,
                  qkv=False,
                  megatron_v2=False,
                  split_qkv=False,
              ):
                  if src_name in sd[0]:
                      dst = getattr(module, dst_name)
                      tmp = sd[0][src_name].cuda()
                      if len(dst.shape) == 1:
                          if split_qkv:
                              dst = mp_replace.qkv_copy(dst, tmp)
                          else:
                              dst = mp_replace.copy(dst, tmp)
                          if qkv and megatron_v2:
                              dst = torch.nn.parameter.Parameter(
                                  _transpose(dst).contiguous()
                              )
                      else:
                          if split_qkv:
                              dst = weight_quantizer.quantize(
                                  mp_replace.qkv_copy(
                                      dst,
                                      tmp
                                      if weight_quantizer.q_int8
                                      else (
                                          transpose(tmp).contiguous()
                                      ),
                                  )
                              )
                          else:
                              dst = weight_quantizer.quantize(
                                  mp_replace.copy(
                                      dst,
                                      tmp
                                      if weight_quantizer.q_int8
                                      else transpose(tmp),
                                  )
                              )
                          if qkv and megatron_v2:
                              scale1 = dst.scale
                              dst = torch.nn.parameter.Parameter(
                                  _transpose(dst).contiguous()
                              )
                              dst.scale = scale1
                      setattr(module, dst_name, dst)
  
              # Extending the maybe_copy function for when the q, k, and v are in separate parameters!
              def maybe_copy_qkv(
                  module, dst_name, src_names, split_qkv=False
              ):
                  if src_names[0] in sd[0]:
                      q = sd[0][src_names[0]]
                      k = sd[0][src_names[1]]
                      v = sd[0][src_names[2]]
                      qkv_data = torch.cat((q, k, v), dim=0)
                      dst = getattr(module, dst_name)
                      if len(dst.shape) == 1:
                          if split_qkv:
                              dst = mp_replace.qkv_copy(
                                  dst, (qkv_data.cuda()).contiguous()
                              )
                          else:
                              dst = mp_replace.copy(
                                  dst, qkv_data.cuda()
                              )
                      else:
                          if split_qkv:
                              dst = weight_quantizer.quantize(
                                  mp_replace.qkv_copy(
                                      dst,
                                      qkv_data.cuda()
                                      if weight_quantizer.q_int8
                                      else (
                                          (
                                              transpose(qkv_data.cuda())
                                          ).contiguous()
                                      ),
                                  )
                              )
                          else:
                              dst = weight_quantizer.quantize(
                                  mp_replace.copy(
                                      dst,
                                      qkv_data.cuda()
                                      if weight_quantizer.q_int8
                                      else transpose(qkv_data.cuda()),
                                  )
                              )
                      setattr(module, dst_name, dst)
  
              if len(param_names) == 14:
                  (
                      qkv_w,
                      qkv_b,
                      attn_ow,
                      attn_ob,
                      mlp_intw,
                      mlp_intb,
                      mlp_ow,
                      mlp_ob,
                      inp_normw,
                      inp_normb,
                      attn_nw,
                      attn_nb,
                      _,
                      split_qkv,
                  ) = param_names
              elif len(param_names) < 14:
                  (
                      q_w,
                      k_w,
                      v_w,
                      attn_ow,
                      mlp_intw,
                      mlp_intb,
                      mlp_ow,
                      mlp_ob,
                      inp_normw,
                      inp_normb,
                      _,
                      split_qkv,
                  ) = param_names
              else:
                  (
                      q_w,
                      q_b,
                      k_w,
                      k_b,
                      v_w,
                      v_b,
                      attn_ow,
                      attn_ob,
                      mlp_intw,
                      mlp_intb,
                      mlp_ow,
                      mlp_ob,
                      inp_normw,
                      inp_normb,
                      attn_nw,
                      attn_nb,
                      _,
                      split_qkv,
                  ) = param_names
  
              maybe_copy(module, "norm_w", prefix + inp_normw)
              maybe_copy(module, "norm_b", prefix + inp_normb)
              if len(param_names) == 14:
                  maybe_copy(
                      module.attention,
                      "attn_qkvw",
                      prefix + qkv_w,
                      qkv=True,
                      megatron_v2=megatron_v2,
                      split_qkv=split_qkv,
                  )
                  maybe_copy(
                      module.attention,
                      "attn_qkvb",
                      prefix + qkv_b,
                      qkv=True,
                      megatron_v2=megatron_v2,
                      split_qkv=split_qkv,
                  )
              elif len(param_names) < 14:
                  maybe_copy_qkv(
                      module.attention,
                      "attn_qkvw",
                      [prefix + q_w, prefix + k_w, prefix + v_w],
                      split_qkv=split_qkv,
                  )
              else:
                  maybe_copy_qkv(
                      module.attention,
                      "attn_qkvw",
                      [prefix + q_w, prefix + k_w, prefix + v_w],
                      split_qkv=split_qkv,
                  )
                  maybe_copy_qkv(
                      module.attention,
                      "attn_qkvb",
                      [prefix + q_b, prefix + k_b, prefix + v_b],
                      split_qkv=split_qkv,
                  )
              maybe_copy(module.attention, "attn_ow", prefix + attn_ow)
              if len(param_names) >= 14:
                  maybe_copy(
                      module.attention, "attn_ob", prefix + attn_ob
                  )
                  maybe_copy(module.mlp, "attn_nw", prefix + attn_nw)
                  maybe_copy(module.mlp, "attn_nb", prefix + attn_nb)
              maybe_copy(module.mlp, "inter_w", prefix + mlp_intw)
              maybe_copy(module.mlp, "inter_b", prefix + mlp_intb)
              maybe_copy(module.mlp, "output_w", prefix + mlp_ow)
              maybe_copy(module.mlp, "output_b", prefix + mlp_ob)
  
      try:
          import transformers
  
          OPTLearnedPositionalEmbedding = (
              transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding
          )
      except:
          OPTLearnedPositionalEmbedding = None
      layer_policies = {
          nn.Linear: load,
          nn.Embedding: load,
          nn.LayerNorm: load,
          EmbeddingLayer: load,
          LinearLayer: load,
          Normalize: load,
          transformer_inference.DeepSpeedTransformerInference: load_transformer_layer,
          OPTLearnedPositionalEmbedding: load,
          OPTEmbedding: load,
      }
  
      all_ds_ids = {}
  
      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 sd[0].keys()
                  ):
                      if hasattr(child, "weight") and (
                          hasattr(child.weight, "ds_id")
                          and child.weight.ds_id in all_ds_ids
                      ):
                          prefix1 = all_ds_ids[child.weight.ds_id]
                          if child.__class__ is nn.Linear:
                              child = LinearLayer(
                                  weight=all_ds_ids[child.weight.ds_id]
                              )
                              setattr(module, name, child)
                      continue
                  child_params = list(child.parameters())
                  if len(child_params) > 0 and (
                      child_params[0].numel() == 0
                      or child_params[0].is_meta
                  ):
                      if child.weight.is_meta:
                          ds_shape = child.weight.shape
                      else:
                          ds_shape = child.weight.ds_shape
                      if child.__class__ is nn.LayerNorm:
                          child = Normalize(
                              dim=ds_shape[-1],
                              dtype=child.weight.dtype,
                              eps=child.eps,
                          )
                          setattr(module, name, child)
                      elif child.__class__ is nn.Linear:
                          child = LinearLayer(
                              weight_shape=child.weight.shape,
                              bias=child.bias,
                          )
                          setattr(module, name, child)
                      elif (
                          child.__class__
                          is OPTLearnedPositionalEmbedding
                      ):
                          child = OPTEmbedding(weight_shape=ds_shape)
                          setattr(module, name, child)
                      else:
                          ds_id = None
                          if hasattr(child.weight, "ds_id"):
                              ds_id = child.weight.ds_id
                          child = EmbeddingLayer(
                              weight_shape=ds_shape,
                              dtype=child.weight.dtype,
                          )
                          if ds_id is not None:
                              all_ds_ids[ds_id] = child.weight
                          setattr(module, name, child)
  
                  layer_policies[child.__class__](
                      child, prefix + name + "."
                  )
              else:
                  load_module_recursive(
                      child,
                      prefix
                      if (level == 0 and ckpt_type == "pp")
                      and param_names[-2]
                      else prefix + name + ".",
                      level + 1,
                  )
  
      load_module_recursive(r_module)
  
      # XXX: hack to tie embedding w. lm_head for BLOOM, need to revist soon
      embedding_weight = None
      for n, p in r_module.named_parameters():
          if "word_embeddings." in n or "embed_tokens." in n:
              embedding_weight = p
  
      if embedding_weight is not None:
          r_module.lm_head.weight = embedding_weight
      for sd_ in sd:
          del sd_
      sd = None
      gc.collect()
  
imports: from torch import nn
  import deepspeed.ops.transformer
  from ..runtime.zero import GatheredParameters
  from .layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding
  import torch
  import gc
label_desc: Use image node_id 1 for calling a global function with example usage: load_model_with_checkpoint(r_module, sd, mp_replace, ckpt_type, weight_quantizer, rank, param_names, transformer_config, megatron_v2) and returns: data, torch, torch