chore: import upstream snapshot with attribution

python/mlc_llm/loader/__init__.py

@@ -0,0 +1,8 @@
+"""
+A subpackage of the compiler that represents mapping between external parameters, quantized
+parameters and parameters in MLC-defined models.
+"""
+
+from .huggingface_loader import HuggingFaceLoader
+from .loader import LOADER, Loader
+from .mapping import ExternMapping, QuantizeMapping

python/mlc_llm/loader/huggingface_loader.py

@@ -0,0 +1,228 @@
+"""A weight loader for HuggingFace's PyTorch format"""
+
+import gc
+import json
+from collections import OrderedDict, defaultdict
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+from tqdm import tqdm
+from tvm.runtime import Device, Tensor
+from tvm.runtime import tensor as as_tensor
+
+from mlc_llm.support import logging
+from mlc_llm.support.preshard import _sharded_param_name
+from mlc_llm.support.style import bold
+
+from .mapping import ExternMapping, QuantizeMapping
+from .stats import Stats
+from .utils import check_parameter_usage, load_safetensor_shard, load_torch_shard
+
+logger = logging.getLogger(__name__)
+
+
+class HuggingFaceLoader:
+    """A loader loading HuggingFace's PyTorch/SafeTensor format and converts them
+    to MLC's parameters.
+
+    Attributes
+    ----------
+    stats : Stats
+        Statistics of the loading process.
+
+    extern_param_map : ExternMapping
+        The parameter mapping from MLC to HuggingFace PyTorch/SafeTensor.
+
+    torch_to_path : Dict[str, Path]
+        A mapping from PyTorch/SafeTensor parameter name to the path of the file containing it,
+        or the path meaning all parameters are stored in a single file.
+
+    cached_files : Dict[Path, Dict[str, np.ndarray]]
+        A cache of the loaded files. The key is the path of the file, and the value is a mapping
+        from parameter name to the parameter value.
+
+    quantize_param_map : Optional[QuantizeMapping]
+        The quantization mapping from MLC to quantized MLC parameters.
+    """
+
+    stats: Stats
+    cached_files: Dict[Path, Dict[str, np.ndarray]]  # noqa: UP006
+    torch_to_path: Dict[str, Path]  # noqa: UP006
+    extern_param_map: ExternMapping
+    quantize_param_map: Optional[QuantizeMapping]
+
+    def __init__(
+        self,
+        path: Path,
+        extern_param_map: ExternMapping,
+        quantize_param_map: Optional[QuantizeMapping] = None,
+    ) -> None:
+        """Create a parameter loader from HuggingFace PyTorch format.
+
+        Parameters
+        ----------
+        path : pathlib.Path
+            Path to either a JSON indexing file, or a PyTorch bin file.
+            1) For JSON indexing file, it is usually `pytorch_model.bin.index.json`
+            or `model.safetensors.index.json` in the repo, which contains a `weight_map` that
+            maps each PyTorch parameter to the file containing the weight.
+            2) For PyTorch bin file, it is usually `pytorch_model.bin` in the repo,
+            which contains all the parameters.
+            3) For safetensor file, it is usually `model.safetensors` in the repo,
+            which contains all the parameters.
+
+        extern_param_map : ExternMapping
+            Maps an MLC parameter to a list of PyTorch/SafeTensor parameters.
+
+        quantize_param_map: Optional[QuantizeMapping]
+            The quantization mapping from MLC to quantized MLC parameters, default to None, which
+            means no quantization.
+        """
+        assert path.is_file(), f"Path {path} is not a file"
+        self.stats = Stats()
+        self.extern_param_map = extern_param_map
+        self.cached_files = {}
+        self.torch_to_path = {}
+        self.quantize_param_map = quantize_param_map
+        if path.suffix in (".bin", ".safetensors", ".pt"):
+            self._load_file(path)
+            for name in self.cached_files[path].keys():
+                self.torch_to_path[name] = path
+        elif path.suffix == ".json":
+            with path.open("r", encoding="utf-8") as in_file:
+                torch_weight_map = json.load(in_file)["weight_map"]
+            for torch_name, path_str in torch_weight_map.items():
+                self.torch_to_path[torch_name] = path.parent / path_str
+        else:
+            raise FileNotFoundError(f"Unknown file suffix: {path}")
+        check_parameter_usage(extern_param_map, set(self.torch_to_path.keys()))
+
+    def load(
+        self,
+        device: Device,
+        preshard_funcs: Optional[Dict[str, Callable]] = None,  # noqa: UP006
+    ) -> Iterator[Tuple[str, Tensor]]:  # noqa: UP006
+        """Load the parameters and yield the MLC parameter and its value.
+
+        Parameters
+        ----------
+        device : Optional[Device]
+            The device to store the parameter, default to None, which means using CPU.
+
+        Yields
+        ------
+        Tuple[str, Tensor]
+            The MLC parameter name and its value, quantized if quantization mapping is provided.
+        """
+        mlc_names = _loading_order(self.extern_param_map, self.torch_to_path)
+        for mlc_name in tqdm(mlc_names):
+            param = self._load_mlc_param(mlc_name, device=device)
+            # Apply quantization if needed, in this case the original parameter may become
+            # multiple quantized parameters.
+            for name, loader_param in self._load_or_quantize(mlc_name, param, device):
+                # Apply presharding if needed
+                if preshard_funcs is not None and name in preshard_funcs:
+                    for shard_id, shard_param in enumerate(preshard_funcs[name](loader_param)):
+                        yield _sharded_param_name(name, shard_id), shard_param
+                else:
+                    yield name, loader_param
+
+        cached_files = list(self.cached_files.keys())
+        for path in cached_files:
+            self._unload_file(path)
+        self.stats.log_time_info("HF")
+        self.stats.log_mem_usage()
+
+    def _load_mlc_param(self, mlc_name: str, device: Optional[Device]) -> Tensor:
+        torch_names = self.extern_param_map.param_map[mlc_name]
+        files_required = {self.torch_to_path[p] for p in torch_names}
+        files_existing = set(self.cached_files.keys())
+        files_to_load = files_required - files_existing
+        files_to_unload = files_existing - files_required
+
+        # Step 1. When there is some file to unloaded:
+        # - If no pending file load: unloading is deferred as there is no gain in peak memory usage;
+        # - Need to load files: unload immediately to save memory and make space for the new files.
+        if files_to_load:
+            for path in files_to_unload:
+                self._unload_file(path)
+        # Step 2. Load all the files needed
+        for path in files_to_load:
+            self._load_file(path)
+        # Step 3. Collect all torch parameters in order
+        torch_params = [self.cached_files[self.torch_to_path[i]][i] for i in torch_names]
+        # Step 4. Apply the mapping function
+        with self.stats.timer("map_time_sec"):
+            param = self.extern_param_map.map_func[mlc_name](*torch_params)
+        if device:
+            return as_tensor(param, device=device)
+        return as_tensor(param)
+
+    def _load_or_quantize(self, mlc_name, param, device: Device):
+        if self.quantize_param_map and mlc_name in self.quantize_param_map.param_map:
+            with self.stats.timer("quant_time_sec"):
+                q_names = self.quantize_param_map.param_map[mlc_name]
+                q_params = self.quantize_param_map.map_func[mlc_name](param)
+                device.sync()
+            for q_name, q_param in zip(q_names, q_params):
+                logger.info(
+                    '[Quantized] Parameter: "%s", shape: %s, dtype: %s',
+                    bold(q_name),
+                    q_param.shape,
+                    q_param.dtype,
+                )
+                yield q_name, q_param
+        else:
+            logger.info(
+                '[Not quantized] Parameter: "%s", shape: %s, dtype: %s',
+                bold(mlc_name),
+                param.shape,
+                param.dtype,
+            )
+            device.sync()
+            yield mlc_name, param
+
+    def _load_file(self, path: Path) -> None:
+        logger.info("Loading HF parameters from: %s", path)
+        load_func = load_safetensor_shard if path.suffix == ".safetensors" else load_torch_shard
+        with self.stats.timer("load_time_sec"):
+            result = {}
+            for name, param in load_func(path):
+                result[name] = param
+                self.stats.mem_add(param.nbytes)
+                if name not in self.extern_param_map.unused_params:
+                    self.stats.total_param_num += param.size
+            self.cached_files[path] = result
+
+    def _unload_file(self, path: Path) -> None:
+        logger.info("Unloading HF weight file: %s", path)
+        with self.stats.timer("load_time_sec"):
+            for _, param in self.cached_files[path].items():
+                self.stats.mem_rm(param.nbytes)
+            del self.cached_files[path]
+            gc.collect()
+
+
+def _loading_order(param_map: ExternMapping, torch_to_path: Dict[str, Path]) -> List[str]:  # noqa: UP006
+    # Step 1. Build a map from path to torch parameters
+    path_to_torch: Dict[Path, List[str]] = defaultdict(list)  # noqa: UP006
+    for torch_name, path in torch_to_path.items():
+        path_to_torch[path].append(torch_name)
+    # Step 2. Build a map from torch parameters to MLC parameters
+    torch_to_mlc = defaultdict(list)
+    for mlc_name, torch_names in param_map.param_map.items():
+        for torch_name in torch_names:
+            torch_to_mlc[torch_name].append(mlc_name)
+    # Step 3. Construct the ordering that ensures file locality
+    order = OrderedDict()
+    for _, torch_names in path_to_torch.items():
+        for torch_name in torch_names:
+            for mlc_name in torch_to_mlc[torch_name]:
+                if mlc_name not in order:
+                    order[mlc_name] = 1
+    return list(order.keys())
+
+
+__all__ = ["HuggingFaceLoader"]

python/mlc_llm/loader/loader.py

@@ -0,0 +1,13 @@
+"""A centralized registry of all existing loaders."""
+
+from typing import Any, Dict  # noqa: UP035
+
+from .huggingface_loader import HuggingFaceLoader
+
+Loader = Any
+
+LOADER: Dict[str, Any] = {  # noqa: UP006
+    "huggingface-torch": HuggingFaceLoader,
+    "huggingface-safetensor": HuggingFaceLoader,
+    "awq": HuggingFaceLoader,
+}

python/mlc_llm/loader/mapping.py

@@ -0,0 +1,102 @@
+"""Parameter mapping for converting different LLM implementations to MLC LLM."""
+
+import dataclasses
+from typing import Callable, Dict, List, Set, Union  # noqa: UP035
+
+import numpy as np
+from tvm.runtime import Tensor
+
+MapFuncVariadic = Union[
+    Callable[[], np.ndarray],
+    Callable[[np.ndarray], np.ndarray],
+    Callable[[np.ndarray, np.ndarray], np.ndarray],
+    Callable[[np.ndarray, np.ndarray, np.ndarray], np.ndarray],
+    Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray],
+]
+
+
+@dataclasses.dataclass
+class ExternMapping:
+    """Mapping from a parameter name in MLC LLM's model definition to its potential source,
+    for example, from MLC parameter "model.layers.2.post_attention_layernorm.weight" to PyTorch's
+    parameter correspondingly.
+
+    Parameters
+    ----------
+    param_map : Dict[str, List[str]]
+        A dictionary that maps the name of a parameter to its source. For example,
+        in Llama2, the source of MLC parameter "model.layers.0.self_attn.qkv_proj.weight" from
+        huggingface torch are:
+
+        - "model.layers.0.self_attn.q_proj.weight"
+        - "model.layers.0.self_attn.k_proj.weight"
+        - "model.layers.0.self_attn.v_proj.weight"
+
+    map_func : Dict[str, Callable[[np.ndarray, ...], np.ndarray]]
+        A dictionary that maps the name of a parameter to a function that combines the source
+        parameters into the MLC parameter. For example, for the above example, the function
+        would be: `lambda q, k, v: np.concatenate([q, k, v], axis=0)`.
+
+    unused_params : Set[str]
+        Parameter names in the source weights that are not used in the MLC LLM model definition.
+    """
+
+    param_map: Dict[str, List[str]] = dataclasses.field(default_factory=dict)  # noqa: UP006
+    map_func: Dict[str, MapFuncVariadic] = dataclasses.field(default_factory=dict)  # noqa: UP006
+    unused_params: Set[str] = dataclasses.field(default_factory=set)  # noqa: UP006
+
+    def add_mapping(
+        self,
+        map_from: str,
+        map_to: List[str],  # noqa: UP006
+        func: MapFuncVariadic,
+    ) -> None:
+        """Add a mapping from MLC parameters to source parametes as well as a mapping function."""
+        self.param_map[map_from] = map_to
+        self.map_func[map_from] = func
+
+    def add_unused(self, name: str):
+        """Add a parameter name in the source parameters to the set of unused parameters."""
+        self.unused_params.add(name)
+
+
+@dataclasses.dataclass
+class QuantizeMapping:
+    """Mapping from a parameter in MLC LLM's model definition to its eventual names and values after
+    quantization. In certain group quantization, for example, `qkv_proj.weight` is mapped to
+    `qkv_proj.weight_quantized` and `qkv_proj.weight_scale` respectively. If a parameter's name is
+    not in the mapping, it is assumed to be unchanged, i.e. not quantized.
+
+    Parameters
+    ----------
+    param_map : Dict[str, List[str]]
+        A dictionary that maps the name of a parameter to its destination. For example,
+        in certain group quantization, the destinations of MLC parameter "qkv_proj.weight` are:
+
+        - "qkv_proj.weight_quantized"
+        - "qkv_proj.weight_scale"
+
+    map_func : Dict[str, Callable[Tensor, List[Tensor]]]
+        A dictionary that maps the name of a parameter to a function that splits the MLC parameter
+        into the destination parameters.
+
+    Notes
+    -----
+    There are two forms of weight conversion in MLC LLM, one is A) on-the-fly quantization to the
+    raw fp16/bf16/fp32 weights from HuggingFace, and the other is B) loading pre-quantized weights
+    from an external framework, e.g. AutoGPTQ, AutoAWQ. From the perspective of parameter
+    correspondence.
+
+    - In case A), it is recommended that the weight loader take both `ExternMapping` and
+    `QuantizeMapping` as input, and do quantiaztion on the fly as a raw parameter being
+    loaded into RAM;
+    - In case B), a pass over `nn.Module` is recommended to take place first to converts parameters
+    from its non-quantized form to the quantized one, and then only `ExternMapping` is
+    used to convert the quantized parameters into the desired form.
+    """
+
+    param_map: Dict[str, List[str]]  # noqa: UP006
+    map_func: Dict[str, Callable[[Tensor], List[Tensor]]]  # noqa: UP006
+
+
+__all__ = ["ExternMapping", "QuantizeMapping"]

python/mlc_llm/loader/standard_loader.py

@@ -0,0 +1,152 @@
+"""Standard HuggingFace loader mapping helpers."""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Iterable, Sequence
+from typing import Callable, Optional, Type  # noqa: UP035
+
+import numpy as np
+from tvm.relax.frontend import nn
+
+from mlc_llm.loader import ExternMapping
+from mlc_llm.quantization import Quantization
+
+NameTransform = Callable[[str], str]
+ExportSpecGetter = Callable[[nn.Module], object]
+
+
+def _default_export_spec(model: nn.Module) -> object:
+    return model.get_default_spec()
+
+
+def make_standard_hf_loader(
+    *,
+    model_cls: Type[nn.Module],  # noqa: UP006
+    layer_prefix: str = "model.layers",
+    qkv_names: Sequence[str] = ("q_proj", "k_proj", "v_proj"),
+    qkv_concat_axis: int = 0,
+    qkv_target_name: str = "qkv_proj",
+    add_qkv_bias: bool = False,
+    qkv_bias_optional: bool = False,
+    gate_up_names: Sequence[str] = ("gate_proj", "up_proj"),
+    gate_up_concat_axis: int = 0,
+    gate_up_target_name: str = "gate_up_proj",
+    include_qkv: bool = True,
+    include_gate_up: bool = True,
+    add_unused: Optional[Iterable[str]] = None,  # noqa: UP045
+    hf_prefix: str = "model.",
+    name_transform: Optional[NameTransform] = None,  # noqa: UP045
+    export_spec_getter: Optional[ExportSpecGetter] = None,  # noqa: UP045
+    num_layers_getter: Optional[Callable[[object], int]] = None,  # noqa: UP045
+) -> Callable[[object, Quantization], ExternMapping]:
+    """Create a standard loader for HuggingFace weights.
+
+    This handles the common QKV concatenation, gate+up concatenation, optional
+    QKV bias mapping, and passes through remaining parameters 1:1.
+    """
+
+    if not qkv_names:
+        include_qkv = False
+    if not gate_up_names:
+        include_gate_up = False
+    if not include_qkv:
+        qkv_names = ()
+    if not include_gate_up:
+        gate_up_names = ()
+
+    def _default_name_transform(name: str) -> str:
+        # When hf_prefix is empty, strip the "model." prefix so models that
+        # expose bare top-level weights (no "model." namespace) still load.
+        if hf_prefix == "":
+            return name[6:] if name.startswith("model.") else name
+        return name
+
+    name_transform_fn = name_transform or _default_name_transform
+    spec_getter = export_spec_getter or _default_export_spec
+    unused_names = tuple(add_unused or ())
+
+    def huggingface(
+        model_config: object,
+        quantization: Quantization,
+    ) -> ExternMapping:
+        model = model_cls(model_config)
+        if quantization is not None:
+            model.to(quantization.model_dtype)
+        _, _named_params, _ = model.export_tvm(
+            spec=spec_getter(model),
+            allow_extern=True,
+        )
+        named_parameters = dict(_named_params)
+        mapping = ExternMapping()
+
+        if include_qkv or include_gate_up or unused_names:
+            if num_layers_getter is None:
+                num_layers = model_config.num_hidden_layers
+            else:
+                num_layers = num_layers_getter(model_config)
+
+            for i in range(num_layers):
+                attn = f"{layer_prefix}.{i}.self_attn"
+                if include_qkv:
+                    mlc_qkv_name = f"{attn}.{qkv_target_name}.weight"
+                    mlc_param = named_parameters[mlc_qkv_name]
+                    mapping.add_mapping(
+                        mlc_qkv_name,
+                        [name_transform_fn(f"{attn}.{name}.weight") for name in qkv_names],
+                        functools.partial(
+                            lambda q, k, v, dtype: np.concatenate(
+                                [q, k, v], axis=qkv_concat_axis
+                            ).astype(dtype),
+                            dtype=mlc_param.dtype,
+                        ),
+                    )
+
+                    if add_qkv_bias:
+                        mlc_bias_name = f"{attn}.{qkv_target_name}.bias"
+                        if (not qkv_bias_optional) or mlc_bias_name in named_parameters:
+                            mlc_param = named_parameters[mlc_bias_name]
+                            mapping.add_mapping(
+                                mlc_bias_name,
+                                [name_transform_fn(f"{attn}.{name}.bias") for name in qkv_names],
+                                functools.partial(
+                                    lambda q, k, v, dtype: np.concatenate(
+                                        [q, k, v], axis=qkv_concat_axis
+                                    ).astype(dtype),
+                                    dtype=mlc_param.dtype,
+                                ),
+                            )
+
+                if include_gate_up:
+                    mlp = f"{layer_prefix}.{i}.mlp"
+                    mlc_gate_up_name = f"{mlp}.{gate_up_target_name}.weight"
+                    if gate_up_names:
+                        mlc_param = named_parameters[mlc_gate_up_name]
+                        mapping.add_mapping(
+                            mlc_gate_up_name,
+                            [name_transform_fn(f"{mlp}.{name}.weight") for name in gate_up_names],
+                            functools.partial(
+                                lambda gate, up, dtype: np.concatenate(
+                                    [gate, up], axis=gate_up_concat_axis
+                                ).astype(dtype),
+                                dtype=mlc_param.dtype,
+                            ),
+                        )
+
+                for unused_name in unused_names:
+                    mapping.add_unused(name_transform_fn(f"{attn}.{unused_name}"))
+
+        for mlc_name, mlc_param in named_parameters.items():
+            if mlc_name not in mapping.param_map:
+                mapping.add_mapping(
+                    mlc_name,
+                    [name_transform_fn(mlc_name)],
+                    functools.partial(
+                        lambda x, dtype: x.astype(dtype),
+                        dtype=mlc_param.dtype,
+                    ),
+                )
+
+        return mapping
+
+    return huggingface

python/mlc_llm/loader/stats.py

@@ -0,0 +1,93 @@
+"""Statistics of the loading process of parameter loaders"""
+
+import dataclasses
+import time
+from contextlib import contextmanager
+
+from mlc_llm.support import logging
+from mlc_llm.support.style import green
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class Stats:
+    """Statistics of the loading process of parameter loaders.
+
+    Attributes
+    ----------
+    load_time_sec : float
+        Time used in loading the parameters.
+
+    map_time_sec : float
+        Time used in applying the mapping function, i.e. `ExternMapping.map_func`.
+
+    quant_time_sec : float
+        Time used in quantizing the parameters, i.e. `QuantizeMapping.quant_func`.
+
+    current_memory_gb : float
+        The current RAM usage in GB.
+
+    total_memory_gb : float
+        The total size data loaded from disk in GB.
+
+    max_memory_gb : float
+        The maximum RAM usage in GB.
+
+    total_param_num: int
+        Total number of parameters (original non-MLC model weights), excluding unused params.
+    """
+
+    load_time_sec: float = 0.0
+    map_time_sec: float = 0.0
+    quant_time_sec: float = 0.0
+
+    current_memory_gb: float = 0.0
+    total_memory_gb: float = 0.0
+    max_memory_gb: float = 0.0
+
+    total_param_num: int = 0
+
+    def timer(self, attr):
+        """A context manager to time the scope and add the time to the attribute."""
+
+        @contextmanager
+        def timed_scope():
+            start_time = time.time()
+            yield
+            elapsed_time = time.time() - start_time
+            setattr(self, attr, getattr(self, attr) + elapsed_time)
+
+        return timed_scope()
+
+    def mem_add(self, nbytes: int):
+        """Add the memory usage by the given number of bytes."""
+        mem_gb = float(nbytes) / float(1024**3)
+        self.current_memory_gb += mem_gb
+        self.total_memory_gb += mem_gb
+        self.max_memory_gb = max(self.max_memory_gb, self.current_memory_gb)
+
+    def mem_rm(self, nbytes: int):
+        """Remove the memory usage by the given number of bytes."""
+        mem_gb = float(nbytes) / float(1024**3)
+        self.current_memory_gb -= mem_gb
+
+    def log_time_info(self, weight_format: str):
+        """Log the time used in loading, pre-quantization and quantization."""
+        logger.info(
+            "%s: %s loading: %.3f sec; Pre-quantization mapping: %.3f sec; Quantization: %.3f sec",
+            green("Time usage"),
+            weight_format,
+            self.load_time_sec,
+            self.map_time_sec,
+            self.quant_time_sec,
+        )
+
+    def log_mem_usage(self):
+        """Log the Memory usage information."""
+        logger.info(
+            "%s: Peak RAM: %.3f GB. Total bytes loaded from disk: %.3f GB",
+            green("RAM usage"),
+            self.max_memory_gb,
+            self.total_memory_gb,
+        )

python/mlc_llm/loader/utils.py

@@ -0,0 +1,75 @@
+"""Common utilities for loading parameters"""
+
+import functools
+import operator
+from collections.abc import Iterator
+from pathlib import Path
+from typing import TYPE_CHECKING, Set, Tuple  # noqa: UP035
+
+import numpy as np
+
+from mlc_llm.support import logging
+
+if TYPE_CHECKING:
+    from .mapping import ExternMapping
+
+
+logger = logging.getLogger(__name__)
+
+
+def check_parameter_usage(param_map: "ExternMapping", extern_weights: Set[str]):  # noqa: UP006
+    """Check that all external parameters have been used and are stored in the weights file."""
+    used_extern_names = set(functools.reduce(operator.iadd, param_map.param_map.values(), []))
+    # Check 1. All extern parameters in the weight files are used unless explicitly specified
+    unused_extern_names = extern_weights - used_extern_names - param_map.unused_params
+    if unused_extern_names:
+        logger.warning(
+            "Unused extern parameters: %s",
+            ", ".join(sorted(unused_extern_names)),
+        )
+    # Check 2. All extern parameters required are stored in the weight files
+    nonexistent_extern_names = used_extern_names - extern_weights
+    if nonexistent_extern_names:
+        raise ValueError(
+            "The following extern parameters do not exist in the weight files:\n  "
+            + "\n  ".join(sorted(nonexistent_extern_names)),
+        )
+
+
+def load_torch_shard(path: Path) -> Iterator[Tuple[str, np.ndarray]]:  # noqa: UP006
+    """Load and yield PyTorch format parameters."""
+    import torch
+
+    for name, param in torch.load(path, map_location=torch.device("cpu")).items():
+        if param is None:
+            logger.warning("Encountered None param, skipping it: %s", name)
+            continue
+        param = param.detach().cpu()
+        dtype = str(param.dtype)
+        if dtype == "torch.bfloat16":
+            param = param.float()
+        param = param.numpy()
+        yield name, param
+
+
+def load_safetensor_shard(path: Path) -> Iterator[Tuple[str, np.ndarray]]:  # noqa: UP006
+    """Load and yield SafeTensor format parameters."""
+    import safetensors
+    import torch
+
+    with safetensors.safe_open(path, framework="pt", device="cpu") as in_file:
+        for name in in_file.keys():
+            param = in_file.get_tensor(name)
+            param = param.detach().cpu()
+            dtype = str(param.dtype)
+            if dtype == "torch.bfloat16":
+                import ml_dtypes
+
+                param = param.view(torch.float16).cpu().numpy().view(ml_dtypes.bfloat16)
+            elif dtype == "torch.float8_e4m3fn":
+                import ml_dtypes
+
+                param = param.view(torch.uint8).cpu().numpy().view(ml_dtypes.float8_e4m3fn)
+            else:
+                param = param.numpy()
+            yield name, param