chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -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