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,250 @@
+"""Python entrypoint of weight conversion."""
+
+import contextlib
+import dataclasses
+import math
+import os
+import tempfile
+from collections.abc import Iterator
+from io import StringIO
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple  # noqa: UP035
+
+from tvm import tirx
+from tvm.contrib import tvmjs
+from tvm.runtime import DataType, Device, Tensor
+from tvm.runtime import cpu as cpu_device
+from tvm.target import Target
+
+from mlc_llm.loader import LOADER
+from mlc_llm.model import Model
+from mlc_llm.quantization import Quantization
+from mlc_llm.support import logging, tqdm
+from mlc_llm.support.auto_weight import detect_weight
+from mlc_llm.support.preshard import apply_preshard
+from mlc_llm.support.style import bold, green
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class ConversionArgs:
+    """Arguments to MLC LLM's weight conversation and quantization flow."""
+
+    config: Path
+    quantization: Quantization
+    model: Model
+    device: Device
+    source: Path
+    source_format: str
+    output: Path
+    lora_adapter: Optional[Path] = None
+
+    def display(self) -> None:
+        """Display the arguments to stdout."""
+
+        def _device_to_str(device: Device) -> str:
+            return f"{Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]}:{device.index}"
+
+        out = StringIO()
+        print(f"{bold('Weight conversion with arguments:')}", file=out)
+        print(f"  {bold('--config'):<25} {self.config}", file=out)
+        print(f"  {bold('--quantization'):<25} {self.quantization}", file=out)
+        print(f"  {bold('--model-type'):<25} {self.model.name}", file=out)
+        print(f"  {bold('--device'):<25} {_device_to_str(self.device)}", file=out)
+        print(f"  {bold('--source'):<25} {self.source}", file=out)
+        print(f"  {bold('--source-format'):<25} {self.source_format}", file=out)
+        print(f"  {bold('--output'):<25} {self.output}", file=out)
+        if self.lora_adapter is not None:
+            print(f"  {bold('--lora-adapter'):<25} {self.lora_adapter}", file=out)
+        print(out.getvalue().rstrip())
+
+
+def _resolve_base_model_dir(source: Path) -> Path:
+    return source if source.is_dir() else source.parent
+
+
+@contextlib.contextmanager
+def _merge_lora_adapter_with_base_model(base_source: Path, lora_adapter: Path) -> Iterator[Path]:
+    base_model_dir = _resolve_base_model_dir(base_source)
+    if not base_model_dir.exists():
+        raise ValueError(f"Base model directory does not exist: {base_model_dir}")
+    if not lora_adapter.exists() or not lora_adapter.is_dir():
+        raise ValueError(f"LoRA adapter directory does not exist: {lora_adapter}")
+
+    try:
+        from peft import PeftModel
+        from transformers import AutoModelForCausalLM
+
+    except ImportError as err:
+        raise ImportError(
+            "`--lora-adapter` requires `peft` and `transformers` to be installed."
+        ) from err
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        merged_model_dir = Path(temp_dir) / "merged_model"
+        logger.info("Merging LoRA adapter %s into base model %s", lora_adapter, base_model_dir)
+
+        base_model = AutoModelForCausalLM.from_pretrained(
+            str(base_model_dir),
+            torch_dtype="auto",
+            trust_remote_code=False,
+            low_cpu_mem_usage=True,
+        )
+        merged_model = PeftModel.from_pretrained(
+            base_model, str(lora_adapter), is_trainable=False
+        ).merge_and_unload()
+        merged_model.save_pretrained(str(merged_model_dir), safe_serialization=True)
+        yield merged_model_dir
+
+
+def _convert_args(args: ConversionArgs) -> None:
+    pre_shards_num = os.getenv("MLC_INTERNAL_PRESHARD_NUM")
+    # model config & quantization config
+    model_config = args.model.config.from_file(args.config)
+    if (
+        args.quantization.kind == "ft-quant"
+        and hasattr(model_config, "tensor_parallel_shards")
+        and model_config.tensor_parallel_shards > 1
+    ):
+        raise NotImplementedError
+    if pre_shards_num is not None:
+        model_config.tensor_parallel_shards = int(pre_shards_num)
+    model, quantize_map = args.model.quantize[args.quantization.kind](
+        model_config, args.quantization
+    )
+    _, _named_params, _ = model.export_tvm(
+        spec=model.get_default_spec(),
+        allow_extern=True,
+    )
+    named_params = dict(_named_params)
+
+    if pre_shards_num is not None:
+        named_params, preshard_funcs = apply_preshard(named_params, int(pre_shards_num), args)
+    else:
+        preshard_funcs = None
+
+    def _check_param(name: str, param: Tensor):
+        nonlocal named_params
+        if name not in named_params:
+            raise ValueError(f"Parameter not found in model: {name}")
+        if name in param_names:
+            raise ValueError(f"Duplication: Parameter {name} already computed")
+
+        # Check shape (possibly dynamic)
+        def _check_shape(actual: tuple, expect: tuple):  # expect can have tirx.Var
+            if len(actual) != len(expect):
+                return False
+            for actual_i, expect_i in zip(actual, expect):
+                assert isinstance(expect_i, (int, tirx.Var))
+                if isinstance(expect_i, int) and actual_i != expect_i:
+                    return False
+            return True
+
+        expect_shape = named_params[name].shape
+        actual_shape = param.shape
+        if not _check_shape(actual_shape, expect_shape):
+            raise ValueError(
+                f"Parameter {name} has shape {param.shape}, but expected {expect_shape}"
+            )
+        # Check dtype
+        actual_dtype = param.dtype
+        expect_dtype = named_params[name].dtype
+        if actual_dtype != expect_dtype:
+            raise ValueError(
+                f"Parameter {name} has dtype {param.dtype}, but expected {expect_dtype}"
+            )
+        del named_params[name]
+
+    # load and quantize
+    param_names = set()
+    total_bytes = 0.0
+    total_params: int = 0
+
+    def _param_generator() -> Iterator[Tuple[str, Tensor]]:  # noqa: UP006
+        nonlocal total_params, total_bytes
+        with Target.from_device(args.device), tqdm.redirect():
+            loader = LOADER[args.source_format](
+                path=args.source,
+                extern_param_map=args.model.source[args.source_format](
+                    model_config, args.quantization
+                ),
+                quantize_param_map=quantize_map,
+            )
+            for name, param in loader.load(device=args.device, preshard_funcs=preshard_funcs):
+                _check_param(name, param)
+                param_names.add(name)
+                param = param.copyto(cpu_device())
+                total_bytes += math.prod(param.shape) * DataType(param.dtype).itemsize
+                yield name, param
+        total_params = loader.stats.total_param_num
+
+    def _metadata_callback() -> Dict[str, Any]:  # noqa: UP006
+        return {
+            "ParamSize": len(param_names),
+            "ParamBytes": total_bytes,
+            "BitsPerParam": total_bytes * 8.0 / total_params,
+        }
+
+    # dump to output directory
+    tvmjs.dump_tensor_cache(
+        _param_generator(),
+        str(args.output),
+        meta_data=_metadata_callback,
+        encode_format="f32-to-bf16",
+        show_progress=False,
+    )
+    if named_params:
+        raise ValueError(f"Parameter not found in source: {', '.join(named_params.keys())}")
+    # Log necessary statistics
+    logger.info(
+        "%s after quantization: %.3f GB",
+        green("Parameter size"),
+        total_bytes / (1024**3),
+    )
+    logger.info(f"%s: {total_params:,}", green("Total parameters"))
+    logger.info(
+        "%s: %.3f",
+        green("Bits per parameter"),
+        total_bytes * 8.0 / total_params,
+    )
+    logger.info("Saved to directory: %s", bold(str(args.output)))
+
+
+def convert_weight(
+    config: Path,
+    quantization: Quantization,
+    model: Model,
+    device: Device,
+    source: Path,
+    source_format: str,
+    output: Path,
+    lora_adapter: Optional[Path] = None,
+):
+    """MLC LLM's weight conversation and quantization flow."""
+    args = ConversionArgs(
+        config, quantization, model, device, source, source_format, output, lora_adapter
+    )
+
+    allowed_lora_source_formats = {"huggingface-safetensor", "huggingface-torch"}
+    if lora_adapter is not None and source_format not in allowed_lora_source_formats:
+        raise ValueError(
+            f"`--lora-adapter` only supports source formats: {sorted(allowed_lora_source_formats)}"
+        )
+
+    if lora_adapter is not None:
+        with _merge_lora_adapter_with_base_model(source, lora_adapter) as merged_model_dir:
+            merged_source, merged_source_format = detect_weight(
+                weight_path=merged_model_dir,
+                config_json_path=config,
+                weight_format="auto",
+            )
+            merged_args = dataclasses.replace(
+                args, source=merged_source, source_format=merged_source_format
+            )
+            merged_args.display()
+            _convert_args(merged_args)
+            return
+
+    args.display()
+    _convert_args(args)