[LoRA] Add an end-to-end single-adapter runtime path for Qwen2

python/mlc_llm/cli/compile.py

@@ -121,6 +121,15 @@ def _check_system_lib_prefix(prefix: str) -> str:
         default=None,
         help=HELP["debug_dump"] + " (default: %(default)s)",
     )
+    parser.add_argument(
+        "--lora-adapter",
+        type=_parse_dir,
+        default=None,
+        help=(
+            "Path to a PEFT LoRA adapter whose separate A/B tensors are compiled into "
+            "the runtime graph (Qwen2/q0f16/TP=1 only)."
+        ),
+    )
     parsed = parser.parse_args(argv)
     target, build_func = detect_target_and_host(
         parsed.device,
@@ -149,4 +158,5 @@ def _check_system_lib_prefix(prefix: str) -> str:
         output=parsed.output,
         overrides=parsed.overrides,
         debug_dump=parsed.debug_dump,
+        lora_adapter=parsed.lora_adapter,
     )

python/mlc_llm/cli/convert_weight.py

@@ -92,6 +92,15 @@ def _parse_lora_adapter(path: Union[str, Path]) -> Path:
             "When provided, adapter weights are merged into the base model before quantization."
         ),
     )
+    parser.add_argument(
+        "--lora-mode",
+        choices=["merge", "runtime"],
+        default="merge",
+        help=(
+            "How to apply --lora-adapter. `merge` folds it into the base weights; "
+            "`runtime` packages separate low-rank tensors (Qwen2/q0f16/TP=1 only)."
+        ),
+    )
 
     parsed = parser.parse_args(argv)
     parsed.source, parsed.source_format = detect_weight(
@@ -109,4 +118,5 @@ def _parse_lora_adapter(path: Union[str, Path]) -> Path:
         source_format=parsed.source_format,
         output=parsed.output,
         lora_adapter=parsed.lora_adapter,
+        lora_mode=parsed.lora_mode,
     )

python/mlc_llm/interface/compile.py

@@ -17,6 +17,7 @@
 from mlc_llm.quantization import Quantization
 from mlc_llm.support import logging
 from mlc_llm.support.config import ConfigBase
+from mlc_llm.support.runtime_lora import RuntimeLoRAConfig, validate_runtime_lora_scope
 from mlc_llm.support.style import bold
 
 from .compiler_flags import ModelConfigOverride, OptimizationFlags
@@ -38,6 +39,7 @@ class CompileArgs:
     output: Path
     overrides: ModelConfigOverride
     debug_dump: Optional[Path]
+    lora_adapter: Optional[Path]
 
     def __post_init__(self) -> None:
         self.opt.update(self.target, self.quantization)
@@ -54,6 +56,8 @@ def display(self) -> None:
         print(f'  {bold("--system-lib-prefix"):<25} "{self.system_lib_prefix}"', file=out)
         print(f"  {bold('--output'):<25} {self.output}", file=out)
         print(f"  {bold('--overrides'):<25} {self.overrides}", file=out)
+        if self.lora_adapter is not None:
+            print(f"  {bold('--lora-adapter'):<25} {self.lora_adapter}", file=out)
         # As it's debug only, no need to display
         # print(f"  {bold('--debug-dump'):<25} {self.debug_dump}", file=out)
         print(out.getvalue().rstrip())
@@ -133,6 +137,13 @@ def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]:  # no
     logger.info("TOP LEVEL MODEL CONFIG BEFORE OVERRIDES: %s", str(model_config))
     _kwargs = getattr(model_config, "kwargs", {})
     model_config = args.overrides.apply(model_config)
+    runtime_lora = getattr(model_config, "runtime_lora", None)
+    if runtime_lora is not None:
+        validate_runtime_lora_scope(
+            model_name=args.model.name,
+            quantization_name=args.quantization.name,
+            tensor_parallel_shards=model_config.tensor_parallel_shards,
+        )
     with args.target:
         op_ext.enable(
             target=args.target,
@@ -192,6 +203,8 @@ def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]:  # no
             "active_vocab_size": avs,
             "model_task": args.model.model_task,
         }
+        if runtime_lora is not None:
+            metadata["runtime_lora"] = runtime_lora.asdict()
         if args.model.embedding_metadata:
             metadata["embedding_metadata"] = dataclasses.asdict(args.model.embedding_metadata)
         logger.info("Registering metadata: %s", metadata)
@@ -236,8 +249,15 @@ def compile(
     output: Path,
     overrides: ModelConfigOverride,
     debug_dump: Optional[Path] = None,
+    lora_adapter: Optional[Path] = None,
 ):
     """Compile a model given its configuration and quantization format to a specific target."""
+    config = dict(config)
+    if "model_config" in config:
+        config["model_config"] = dict(config["model_config"])
+    if lora_adapter is not None:
+        config["runtime_lora"] = RuntimeLoRAConfig.from_peft_directory(lora_adapter).asdict()
+
     avs = None
     if "active_vocab_size" in config:
         avs = config.pop("active_vocab_size")
@@ -260,6 +280,7 @@ def compile(
         output,
         overrides,
         debug_dump,
+        lora_adapter,
     )
     args.display()
     _compile(args, model_config)

python/mlc_llm/interface/convert_weight.py

@@ -17,11 +17,16 @@
 from tvm.target import Target
 
 from mlc_llm.loader import LOADER
+from mlc_llm.loader.runtime_lora import (
+    make_runtime_lora_mapping,
+    resolve_runtime_lora_weight,
+)
 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.runtime_lora import RuntimeLoRAConfig, validate_runtime_lora_scope
 from mlc_llm.support.style import bold, green
 
 logger = logging.getLogger(__name__)
@@ -39,6 +44,8 @@ class ConversionArgs:
     source_format: str
     output: Path
     lora_adapter: Optional[Path] = None
+    lora_mode: str = "merge"
+    runtime_lora: Optional[RuntimeLoRAConfig] = None
 
     def display(self) -> None:
         """Display the arguments to stdout."""
@@ -57,6 +64,7 @@ def _device_to_str(device: Device) -> str:
         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(f"  {bold('--lora-mode'):<25} {self.lora_mode}", file=out)
         print(out.getvalue().rstrip())
 
 
@@ -102,6 +110,13 @@ 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.runtime_lora is not None:
+        model_config.runtime_lora = args.runtime_lora
+        validate_runtime_lora_scope(
+            model_name=args.model.name,
+            quantization_name=args.quantization.name,
+            tensor_parallel_shards=model_config.tensor_parallel_shards,
+        )
     if (
         args.quantization.kind == "ft-quant"
         and hasattr(model_config, "tensor_parallel_shards")
@@ -163,6 +178,7 @@ def _check_shape(actual: tuple, expect: tuple):  # expect can have tirx.Var
 
     def _param_generator() -> Iterator[Tuple[str, Tensor]]:  # noqa: UP006
         nonlocal total_params, total_bytes
+        loaders = []
         with Target.from_device(args.device), tqdm.redirect():
             loader = LOADER[args.source_format](
                 path=args.source,
@@ -171,13 +187,30 @@ def _param_generator() -> Iterator[Tuple[str, Tensor]]:  # noqa: UP006
                 ),
                 quantize_param_map=quantize_map,
             )
+            loaders.append(loader)
             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
+
+            if args.runtime_lora is not None:
+                assert args.lora_adapter is not None
+                adapter_weight = resolve_runtime_lora_weight(args.lora_adapter)
+                adapter_loader = LOADER["huggingface-safetensor"](
+                    path=adapter_weight,
+                    extern_param_map=make_runtime_lora_mapping(named_params, adapter_weight),
+                    quantize_param_map=None,
+                )
+                loaders.append(adapter_loader)
+                for name, param in adapter_loader.load(device=args.device):
+                    _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 = sum(item.stats.total_param_num for item in loaders)
 
     def _metadata_callback() -> Dict[str, Any]:  # noqa: UP006
         return {
@@ -220,10 +253,24 @@ def convert_weight(
     source_format: str,
     output: Path,
     lora_adapter: Optional[Path] = None,
+    lora_mode: str = "merge",
 ):
     """MLC LLM's weight conversation and quantization flow."""
+    if lora_mode not in ("merge", "runtime"):
+        raise ValueError(f"Unknown LoRA conversion mode: {lora_mode}")
+    if lora_adapter is None and lora_mode != "merge":
+        raise ValueError("`--lora-mode runtime` requires `--lora-adapter`.")
+
     args = ConversionArgs(
-        config, quantization, model, device, source, source_format, output, lora_adapter
+        config,
+        quantization,
+        model,
+        device,
+        source,
+        source_format,
+        output,
+        lora_adapter,
+        lora_mode,
     )
 
     allowed_lora_source_formats = {"huggingface-safetensor", "huggingface-torch"}
@@ -232,6 +279,13 @@ def convert_weight(
             f"`--lora-adapter` only supports source formats: {sorted(allowed_lora_source_formats)}"
         )
 
+    if lora_adapter is not None and lora_mode == "runtime":
+        runtime_lora = RuntimeLoRAConfig.from_peft_directory(lora_adapter)
+        runtime_args = dataclasses.replace(args, runtime_lora=runtime_lora)
+        runtime_args.display()
+        _convert_args(runtime_args)
+        return
+
     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(

python/mlc_llm/loader/runtime_lora.py

@@ -0,0 +1,96 @@
+"""Loader mapping for PEFT weights used by runtime LoRA branches."""
+
+from __future__ import annotations
+
+import functools
+import json
+import re
+from pathlib import Path
+from typing import Dict, Iterable  # noqa: UP035
+
+from tvm.relax.frontend import nn
+
+from .mapping import ExternMapping
+
+
+def _cast_runtime_lora_weight(tensor, *, dtype, scale: float):
+    if scale != 1.0:
+        tensor = tensor * scale
+    return tensor.astype(dtype)
+
+
+def resolve_runtime_lora_weight(adapter_dir: Path) -> Path:
+    """Resolve the safetensors file (or index) in a PEFT adapter directory."""
+    for filename in (
+        "adapter_model.safetensors.index.json",
+        "adapter_model.safetensors",
+    ):
+        candidate = adapter_dir / filename
+        if candidate.is_file():
+            return candidate
+    if (adapter_dir / "adapter_model.bin").is_file():
+        raise ValueError(
+            "Runtime LoRA currently requires safetensors adapter weights; "
+            "`adapter_model.bin` is not supported."
+        )
+    raise ValueError(f"Cannot find PEFT safetensors weights in: {adapter_dir}")
+
+
+def _read_safetensor_keys(path: Path) -> Iterable[str]:
+    if path.suffix == ".json":
+        with path.open("r", encoding="utf-8") as index_file:
+            return tuple(json.load(index_file)["weight_map"].keys())
+
+    with path.open("rb") as weight_file:
+        header_length_bytes = weight_file.read(8)
+        if len(header_length_bytes) != 8:
+            raise ValueError(f"Invalid safetensors header: {path}")
+        header_length = int.from_bytes(header_length_bytes, byteorder="little", signed=False)
+        header = json.loads(weight_file.read(header_length).decode("utf-8"))
+    return tuple(name for name in header if name != "__metadata__")
+
+
+def _normalize_peft_name(name: str) -> str:
+    return re.sub(r"\.(lora_[AB])\.[^.]+\.weight$", r".\1.weight", name)
+
+
+def make_runtime_lora_mapping(
+    named_parameters: Dict[str, nn.Parameter],  # noqa: UP006
+    adapter_weight: Path,
+) -> ExternMapping:
+    """Map exported runtime-LoRA parameters to their PEFT tensor names."""
+    adapter_keys = tuple(_read_safetensor_keys(adapter_weight))
+    normalized_keys = {name: _normalize_peft_name(name) for name in adapter_keys}
+    mapping = ExternMapping()
+
+    for mlc_name, mlc_param in named_parameters.items():
+        source_suffix = mlc_param.attrs.get("peft_source_suffix")
+        if source_suffix is None:
+            continue
+        normalized_suffix = _normalize_peft_name(source_suffix)
+        candidates = [
+            name
+            for name, normalized in normalized_keys.items()
+            if normalized.endswith(normalized_suffix)
+        ]
+        if len(candidates) != 1:
+            raise ValueError(
+                f"Expected one PEFT tensor ending in `{source_suffix}` for `{mlc_name}`, "
+                f"but found {len(candidates)}."
+            )
+        mapping.add_mapping(
+            mlc_name,
+            candidates,
+            functools.partial(
+                _cast_runtime_lora_weight,
+                dtype=mlc_param.dtype,
+                scale=float(mlc_param.attrs.get("runtime_lora_scale", 1.0)),
+            ),
+        )
+
+    if not mapping.param_map:
+        raise ValueError("The exported model does not contain runtime LoRA parameters.")
+    return mapping
+
+
+__all__: tuple[str, ...] = ("make_runtime_lora_mapping", "resolve_runtime_lora_weight")

python/mlc_llm/loader/standard_loader.py

@@ -137,6 +137,8 @@ def huggingface(
                     mapping.add_unused(name_transform_fn(f"{attn}.{unused_name}"))
 
         for mlc_name, mlc_param in named_parameters.items():
+            if mlc_param.attrs.get("runtime_lora_param", False):
+                continue
             if mlc_name not in mapping.param_map:
                 mapping.add_mapping(
                     mlc_name,