docstr: Load checkpoint weights.
  
  # Args.
      checkpoint_path: str or Path,
          a checkpoint directory where FT checkpoint files locate.
      weight_dtype: str or np.dtype, the data type of stored weights.
docstr_tok: ["Load","checkpoint","weights",".","#","Args",".","checkpoint_path",":","str","or","Path",",","a","checkpoint","directory","where","FT","checkpoint","files","locate",".","weight_dtype",":","str","or","np.dtype",",","the","data","type","of","stored","weights","."]
code: def load(
      self,
      checkpoint_path: PathLike,
      compute_dtype: torch.dtype,
      weight_dtype: Optional[Union[str, np.dtype]] = None,
      device: Optional[Union[int, str, torch.device]] = None,
  ):
      """Load checkpoint weights.
  
      # Args.
          checkpoint_path: str or Path,
              a checkpoint directory where FT checkpoint files locate.
          weight_dtype: str or np.dtype, the data type of stored weights.
      """
  
      checkpoint_path = Path(checkpoint_path)
      if not checkpoint_path.exists():
          raise FileNotFoundError(
              f"Could not find checkpoint {str(checkpoint_path)}"
          )
  
      weight_dtype = to_numpy_dtype(weight_dtype)
      print(
          f"Load weights from {str(checkpoint_path)} (data type: {weight_dtype}"  # noqa: E501
      )
  
      self.weights = list()
      self.int8_weights = list()
      self.int8_scales = list()
      torch.cuda.empty_cache()
  
      def _load_from_file(fname):
          quant_sub_names = [
              "attention.query_key_value.weight",
              "attention.dense.weight",
              "dense_h_to_4h.weight",
              "dense_4h_to_h.weight",
          ]
          _weight = torch.from_numpy(
              np.fromfile(checkpoint_path / fname, dtype=weight_dtype)
          )
          _weight = _weight.to(compute_dtype)
          weight_index = len(self.weights)
          expected_shape = self.expected_weight_shapes[weight_index]
  
          try:
              if _weight.nelement() > 0:
                  _weight = _weight.reshape(expected_shape)
          except:  # noqa: E722
              raise ValueError(
                  f"num_heads, size_per_head, vocab_size, and max_seq_len must be the same "  # noqa: E501
                  f"as the ones during training (weight: {fname} expected shape: {expected_shape}, "  # noqa: E501
                  f"got shape: {_weight.shape})."
              )
  
          should_quantize = any(
              sub_name in fname for sub_name in quant_sub_names
          )
          if self.int8_mode != 0 and should_quantize:
              calibrate = self.weight_transpose_calibrate_quantize
              int8_weight, int8_scales = calibrate(_weight)
  
              # int8 weights should appear in same order as FP weights.
              # Move to device and add to the int8 list.
              dummy_weight = torch.empty(0, dtype=compute_dtype)
              if device is not None:
                  int8_weight = int8_weight.to(device)
                  int8_scales = int8_scales.to(device)
                  dummy_weight = dummy_weight.to(device)
  
              self.int8_weights.append(int8_weight)
              self.int8_scales.append(int8_scales)
              self.weights.append(dummy_weight)
          else:
              if device is not None:
                  _weight = _weight.to(device)
              self.weights.append(_weight)
  
      # Load
      # pylint:disable=line-too-long
      layer_offset = self.local_num_layers * self.pipeline_para_rank
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.input_layernorm.weight.bin"
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.input_layernorm.bias.bin"
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.attention.query_key_value.weight.{self.tensor_para_rank}.bin"  # noqa: E501
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.attention.query_key_value.bias.{self.tensor_para_rank}.bin"  # noqa: E501
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.attention.dense.weight.{self.tensor_para_rank}.bin"  # noqa: E501
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.attention.dense.bias.bin"
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.post_attention_layernorm.weight.bin"  # noqa: E501
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.post_attention_layernorm.bias.bin"  # noqa: E501
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.mlp.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.mlp.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.mlp.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
          )
          for i in range(self.local_num_layers)
      ]
      [
          _load_from_file(
              f"model.layers.{layer_offset + i}.mlp.dense_4h_to_h.bias.bin"
          )
          for i in range(self.local_num_layers)
      ]
  
      if self.has_adapters:
          [
              _load_from_file(
                  f"model.layers.{layer_offset + i}.after_attention_adapter.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
              )
              for i in range(self.local_num_layers)
          ]
          [
              _load_from_file(
                  f"model.layers.{layer_offset + i}.after_attention_adapter.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
              )
              for i in range(self.local_num_layers)
          ]
          [
              _load_from_file(
                  f"model.layers.{layer_offset + i}.after_attention_adapter.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
              )
              for i in range(self.local_num_layers)
          ]
          [
              _load_from_file(
                  f"model.layers.{layer_offset + i}.after_attention_adapter.dense_4h_to_h.bias.bin"  # noqa: E501
              )
              for i in range(self.local_num_layers)
          ]
          [
              _load_from_file(
                  f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
              )
              for i in range(self.local_num_layers)
          ]
          [
              _load_from_file(
                  f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
              )
              for i in range(self.local_num_layers)
          ]
          [
              _load_from_file(
                  f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
              )
              for i in range(self.local_num_layers)
          ]
          [
              _load_from_file(
                  f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_4h_to_h.bias.bin"  # noqa: E501
              )
              for i in range(self.local_num_layers)
          ]
  
      assert len(self.weights) == len(
          self.expected_weight_shapes
      ), "Incorrect number of weights loaded"
  
imports: import math
  from abc import abstractmethod
  from pathlib import Path
  from typing import List, Literal, Optional, Union
  import os
  import numpy
  from .None import comm
  from .None import profiler
  from .gpt import GptInitModelParameters
  from nebullvm.optional_modules.torch import torch
label_desc: Use image node_id 13 for calling the GptLayerWeights obj's underlying member method code with example usage: obj.load(checkpoint_path, compute_dtype, weight_dtype, device) without return types