chore: import upstream snapshot with attribution

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