code: def load(module, state_dict, prefix, mp_group=None):
      mp_replace = ReplaceWithTensorSlicing(mp_group=mp_group)
      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 = mp_replace.strided_copy(
                      module.weight.data,
                      state_dict[prefix + "weight"],
                      num_splits=3,
                  )
              else:
                  module.weight = mp_replace.copy(
                      module.weight.data, state_dict[prefix + "weight"]
                  )
      else:
          if hasattr(module, "norm") and hasattr(module.norm, "weight"):
              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 = mp_replace.copy(
                  module.norm.weight.data, state_dict[prefix + "weight"]
              )
  
      if prefix + "bias" in state_dict.keys():
          if hasattr(module, "bias"):
              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,
                  )
              module.bias = mp_replace.copy(
                  module.bias, state_dict[prefix + "bias"]
              )
          else:
              if hasattr(module, "norm") and hasattr(
                  module.norm, "bias"
              ):
                  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 = mp_replace.copy(
                      module.norm.bias, state_dict[prefix + "bias"]
                  )
  
imports: import re
  from torch import nn
  from .replace_policy import replace_policies
  from typing import Optional
  import torch
  from deepspeed import comm
  from .layers import LinearAllreduce, LinearLayer, LmHeadLinearAllreduce
  from deepspeed.accelerator import get_accelerator
  from .fusedqkv_utils import require_tp_fused_qkvw, prepare_tp_fused_qkvw
  from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
label_desc: Use image node_id 3 for calling the Loading obj's underlying member method code with example usage: obj.load(module, state_dict, prefix, mp_group) without return types