From 5d57129e966330c99c5e0df393a202d4314eaf8e Mon Sep 17 00:00:00 2001 From: yash solanki Date: Fri, 17 Jul 2026 13:19:52 +0530 Subject: [PATCH 1/2] [LoRA] Add Qwen2 single-adapter runtime path --- python/mlc_llm/cli/compile.py | 10 ++ python/mlc_llm/cli/convert_weight.py | 10 ++ python/mlc_llm/interface/compile.py | 21 +++ python/mlc_llm/interface/convert_weight.py | 58 ++++++- python/mlc_llm/loader/runtime_lora.py | 96 +++++++++++ python/mlc_llm/loader/standard_loader.py | 2 + python/mlc_llm/model/qwen2/qwen2_model.py | 116 ++++++++++++- python/mlc_llm/support/runtime_lora.py | 163 ++++++++++++++++++ .../python/support/test_cli_convert_weight.py | 51 ++++++ .../support/test_convert_weight_lora_merge.py | 46 +++++ .../python/support/test_qwen2_runtime_lora.py | 65 +++++++ tests/python/support/test_runtime_lora.py | 158 +++++++++++++++++ 12 files changed, 787 insertions(+), 9 deletions(-) create mode 100644 python/mlc_llm/loader/runtime_lora.py create mode 100644 python/mlc_llm/support/runtime_lora.py create mode 100644 tests/python/support/test_qwen2_runtime_lora.py create mode 100644 tests/python/support/test_runtime_lora.py diff --git a/python/mlc_llm/cli/compile.py b/python/mlc_llm/cli/compile.py index 233160c31b..a7aeaabc1c 100644 --- a/python/mlc_llm/cli/compile.py +++ b/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, ) diff --git a/python/mlc_llm/cli/convert_weight.py b/python/mlc_llm/cli/convert_weight.py index b489f3e34d..445b094655 100644 --- a/python/mlc_llm/cli/convert_weight.py +++ b/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, ) diff --git a/python/mlc_llm/interface/compile.py b/python/mlc_llm/interface/compile.py index 053d3afc96..e7d88eda27 100644 --- a/python/mlc_llm/interface/compile.py +++ b/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) diff --git a/python/mlc_llm/interface/convert_weight.py b/python/mlc_llm/interface/convert_weight.py index c5f3325011..2136e517fa 100644 --- a/python/mlc_llm/interface/convert_weight.py +++ b/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( diff --git a/python/mlc_llm/loader/runtime_lora.py b/python/mlc_llm/loader/runtime_lora.py new file mode 100644 index 0000000000..41a50dd303 --- /dev/null +++ b/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") diff --git a/python/mlc_llm/loader/standard_loader.py b/python/mlc_llm/loader/standard_loader.py index 5444c1e6c8..ff6341d390 100644 --- a/python/mlc_llm/loader/standard_loader.py +++ b/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, diff --git a/python/mlc_llm/model/qwen2/qwen2_model.py b/python/mlc_llm/model/qwen2/qwen2_model.py index 71eb12e839..3e620713b4 100644 --- a/python/mlc_llm/model/qwen2/qwen2_model.py +++ b/python/mlc_llm/model/qwen2/qwen2_model.py @@ -16,6 +16,7 @@ from mlc_llm.support import logging from mlc_llm.support import tensor_parallel as tp from mlc_llm.support.config import ConfigBase +from mlc_llm.support.runtime_lora import RuntimeLoRAConfig from mlc_llm.support.style import bold logger = logging.getLogger(__name__) @@ -41,9 +42,14 @@ class QWen2Config(ConfigBase): head_dim: int = 0 dtype: str = "float32" max_batch_size: int = 1 + runtime_lora: Optional[RuntimeLoRAConfig] = None kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 def __post_init__(self): + if isinstance(self.runtime_lora, dict): + self.runtime_lora = RuntimeLoRAConfig.from_dict(self.runtime_lora) + if self.runtime_lora is not None and self.tensor_parallel_shards != 1: + raise ValueError("Runtime LoRA currently requires `tensor_parallel_shards=1`.") if self.context_window_size == 0: for name in ["max_position_embeddings", "max_sequence_length"]: if name in self.kwargs: @@ -81,8 +87,44 @@ def __post_init__(self): self.prefill_chunk_size = min(self.context_window_size, 8192) +class RuntimeLoRALinear(nn.Module): + """An inference-only low-rank branch backed by PEFT A and B tensors.""" + + def __init__( + self, + in_features: int, + out_features: int, + source_module: str, + config: RuntimeLoRAConfig, + ): + self.lora_a = nn.Linear(in_features, config.rank, bias=False) + self.lora_b = nn.Linear(config.rank, out_features, bias=False) + self.lora_a.weight.attrs["runtime_lora_param"] = True + self.lora_b.weight.attrs["runtime_lora_param"] = True + self.lora_a.weight.attrs["peft_source_suffix"] = f"{source_module}.lora_A.weight" + self.lora_b.weight.attrs["peft_source_suffix"] = f"{source_module}.lora_B.weight" + # Fold the constant PEFT scale into B while loading to avoid per-token scaling. + self.lora_b.weight.attrs["runtime_lora_scale"] = config.scaling + + def forward(self, x: Tensor): + return self.lora_b(self.lora_a(x)) + + +def _make_runtime_lora( + config: QWen2Config, + target_module: str, + in_features: int, + out_features: int, + source_module: str, +) -> Optional[RuntimeLoRALinear]: + runtime_lora = config.runtime_lora + if runtime_lora is None or not runtime_lora.applies_to(target_module): + return None + return RuntimeLoRALinear(in_features, out_features, source_module, runtime_lora) + + class QWen2Attention(nn.Module): - def __init__(self, config: QWen2Config): + def __init__(self, config: QWen2Config, layer_id: int): self.head_dim = config.head_dim if config.num_key_value_heads % config.tensor_parallel_shards != 0: raise ValueError( @@ -101,11 +143,39 @@ def __init__(self, config: QWen2Config): self.o_proj = nn.Linear( self.num_attention_heads * self.head_dim, config.hidden_size, bias=False ) + prefix = f"model.layers.{layer_id}.self_attn" + q_size = self.num_attention_heads * self.head_dim + kv_size = self.num_key_value_heads * self.head_dim + self.q_proj_lora = _make_runtime_lora( + config, "q_proj", config.hidden_size, q_size, f"{prefix}.q_proj" + ) + self.k_proj_lora = _make_runtime_lora( + config, "k_proj", config.hidden_size, kv_size, f"{prefix}.k_proj" + ) + self.v_proj_lora = _make_runtime_lora( + config, "v_proj", config.hidden_size, kv_size, f"{prefix}.v_proj" + ) + self.o_proj_lora = _make_runtime_lora( + config, "o_proj", q_size, config.hidden_size, f"{prefix}.o_proj" + ) def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): d, h_q, h_kv = self.head_dim, self.num_attention_heads, self.num_key_value_heads b, s, _ = hidden_states.shape qkv = self.c_attn(hidden_states) + if any( + branch is not None for branch in (self.q_proj_lora, self.k_proj_lora, self.v_proj_lora) + ): + q_size = h_q * d + kv_size = h_kv * d + q, k, v = op.split(qkv, [q_size, q_size + kv_size], axis=-1) + if self.q_proj_lora is not None: + q = q + self.q_proj_lora(hidden_states) + if self.k_proj_lora is not None: + k = k + self.k_proj_lora(hidden_states) + if self.v_proj_lora is not None: + v = v + self.v_proj_lora(hidden_states) + qkv = op.concat([q, k, v], dim=-1) qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) output = op.reshape( paged_kv_cache.attention_with_fused_qkv( @@ -114,6 +184,8 @@ def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: (b, s, h_q * d), ) attn_output = self.o_proj(output) + if self.o_proj_lora is not None: + attn_output = attn_output + self.o_proj_lora(output) return attn_output @@ -140,7 +212,7 @@ def lm_head_forward(self, x: nn.Tensor): class QWen2MLP(nn.Module): - def __init__(self, config: QWen2Config): + def __init__(self, config: QWen2Config, layer_id: int): if config.intermediate_size % config.tensor_parallel_shards != 0: raise ValueError( f"Cannot split MLP intermediate size {config.intermediate_size} " @@ -150,17 +222,47 @@ def __init__(self, config: QWen2Config): self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] + prefix = f"model.layers.{layer_id}.mlp" + self.gate_proj_lora = _make_runtime_lora( + config, + "gate_proj", + config.hidden_size, + self.intermediate_size, + f"{prefix}.gate_proj", + ) + self.up_proj_lora = _make_runtime_lora( + config, + "up_proj", + config.hidden_size, + self.intermediate_size, + f"{prefix}.up_proj", + ) + self.down_proj_lora = _make_runtime_lora( + config, + "down_proj", + self.intermediate_size, + config.hidden_size, + f"{prefix}.down_proj", + ) def forward(self, x: Tensor): concat_x1_x2 = self.gate_up_proj(x) x1, x2 = op.split(concat_x1_x2, 2, axis=-1) - return self.down_proj(self.act_fn(x1) * x2) + if self.gate_proj_lora is not None: + x1 = x1 + self.gate_proj_lora(x) + if self.up_proj_lora is not None: + x2 = x2 + self.up_proj_lora(x) + hidden_states = self.act_fn(x1) * x2 + output = self.down_proj(hidden_states) + if self.down_proj_lora is not None: + output = output + self.down_proj_lora(hidden_states) + return output class QWen2DecoderLayer(nn.Module): - def __init__(self, config: QWen2Config): - self.self_attn = QWen2Attention(config) - self.mlp = QWen2MLP(config) + def __init__(self, config: QWen2Config, layer_id: int): + self.self_attn = QWen2Attention(config, layer_id) + self.mlp = QWen2MLP(config, layer_id) self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) self.post_attention_layernorm = nn.RMSNorm( config.hidden_size, -1, config.rms_norm_eps, bias=False @@ -212,7 +314,7 @@ class QWen2Model(nn.Module): def __init__(self, config: QWen2Config): self.embed_tokens = Qwen2Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList( - [QWen2DecoderLayer(config) for _ in range(config.num_hidden_layers)] + [QWen2DecoderLayer(config, layer_id) for layer_id in range(config.num_hidden_layers)] ) self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) diff --git a/python/mlc_llm/support/runtime_lora.py b/python/mlc_llm/support/runtime_lora.py new file mode 100644 index 0000000000..9e2babf2d5 --- /dev/null +++ b/python/mlc_llm/support/runtime_lora.py @@ -0,0 +1,163 @@ +"""Validation and normalized metadata for the initial runtime LoRA path.""" + +from __future__ import annotations + +import dataclasses +import json +import math +from pathlib import Path +from typing import Any, Dict, Tuple # noqa: UP035 + +SUPPORTED_QWEN2_TARGET_MODULES = frozenset( + { + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + } +) + + +@dataclasses.dataclass(frozen=True) +class RuntimeLoRAConfig: + """The subset of PEFT LoRA metadata supported by runtime LoRA.""" + + rank: int + alpha: float + target_modules: Tuple[str, ...] # noqa: UP006 + use_rslora: bool = False + + @property + def scaling(self) -> float: + """Return the inference-time LoRA scale.""" + denominator = math.sqrt(self.rank) if self.use_rslora else self.rank + return self.alpha / denominator + + def applies_to(self, module_name: str) -> bool: + """Return whether the adapter targets a logical projection.""" + return module_name in self.target_modules + + def asdict(self) -> Dict[str, Any]: # noqa: UP006 + """Return JSON-compatible normalized metadata.""" + return { + "rank": self.rank, + "alpha": self.alpha, + "target_modules": list(self.target_modules), + "use_rslora": self.use_rslora, + } + + @classmethod + def from_dict(cls, source: Dict[str, Any]) -> RuntimeLoRAConfig: # noqa: UP006 + """Load normalized runtime metadata.""" + return cls._create( + rank=source.get("rank"), + alpha=source.get("alpha"), + target_modules=source.get("target_modules"), + use_rslora=source.get("use_rslora", False), + ) + + @classmethod + def from_peft_directory(cls, adapter_dir: Path) -> RuntimeLoRAConfig: + """Load and strictly validate a PEFT ``adapter_config.json``.""" + config_path = adapter_dir / "adapter_config.json" + if not config_path.is_file(): + raise ValueError(f"PEFT adapter config does not exist: {config_path}") + with config_path.open("r", encoding="utf-8") as config_file: + source = json.load(config_file) + if not isinstance(source, dict): + raise ValueError("PEFT adapter config must be a JSON object.") + + if str(source.get("peft_type", "")).upper() != "LORA": + raise ValueError('Runtime LoRA requires a PEFT adapter with `peft_type="LORA"`.') + task_type = source.get("task_type") + if task_type not in (None, "CAUSAL_LM"): + raise ValueError( + "Runtime LoRA currently supports only PEFT causal-language-model adapters." + ) + + unsupported_options = { + "alora_invocation_tokens": source.get("alora_invocation_tokens"), + "alpha_pattern": source.get("alpha_pattern"), + "arrow_config": source.get("arrow_config"), + "ensure_weight_tying": source.get("ensure_weight_tying"), + "exclude_modules": source.get("exclude_modules"), + "layer_replication": source.get("layer_replication"), + "layers_pattern": source.get("layers_pattern"), + "layers_to_transform": source.get("layers_to_transform"), + "lora_bias": source.get("lora_bias"), + "megatron_config": source.get("megatron_config"), + "monteclora_config": source.get("monteclora_config"), + "modules_to_save": source.get("modules_to_save"), + "rank_pattern": source.get("rank_pattern"), + "target_parameters": source.get("target_parameters"), + "trainable_token_indices": source.get("trainable_token_indices"), + "use_bdlora": source.get("use_bdlora"), + "use_qalora": source.get("use_qalora"), + } + enabled_options = [name for name, value in unsupported_options.items() if value] + if enabled_options: + raise ValueError( + "Runtime LoRA does not yet support PEFT options: " + ", ".join(enabled_options) + ) + if source.get("bias", "none") != "none": + raise ValueError('Runtime LoRA currently requires PEFT `bias="none"`.') + if source.get("fan_in_fan_out", False): + raise ValueError("Runtime LoRA does not support `fan_in_fan_out=True`.") + if source.get("use_dora", False): + raise ValueError("Runtime LoRA does not support DoRA adapters.") + + return cls._create( + rank=source.get("r"), + alpha=source.get("lora_alpha"), + target_modules=source.get("target_modules"), + use_rslora=source.get("use_rslora", False), + ) + + @classmethod + def _create(cls, rank, alpha, target_modules, use_rslora) -> RuntimeLoRAConfig: + if not isinstance(rank, int) or isinstance(rank, bool) or rank <= 0: + raise ValueError("Runtime LoRA requires a positive integer rank.") + if not isinstance(alpha, (int, float)) or isinstance(alpha, bool): + raise ValueError("Runtime LoRA requires a numeric `lora_alpha`.") + alpha = float(alpha) + if not math.isfinite(alpha): + raise ValueError("Runtime LoRA requires a finite `lora_alpha`.") + if isinstance(target_modules, str) or not isinstance(target_modules, (list, tuple, set)): + raise ValueError("Runtime LoRA requires PEFT `target_modules` to be an explicit list.") + if not all(isinstance(target, str) for target in target_modules): + raise ValueError("Runtime LoRA requires every PEFT target module to be a string.") + if not isinstance(use_rslora, bool): + raise ValueError("Runtime LoRA requires PEFT `use_rslora` to be a boolean.") + targets = tuple(sorted(set(target_modules))) + if not targets: + raise ValueError("Runtime LoRA requires at least one target module.") + unsupported_targets = sorted(set(targets) - SUPPORTED_QWEN2_TARGET_MODULES) + if unsupported_targets: + raise ValueError( + "Runtime LoRA currently supports only Qwen2 projection targets; unsupported: " + + ", ".join(unsupported_targets) + ) + return cls( + rank=rank, + alpha=alpha, + target_modules=targets, + use_rslora=use_rslora, + ) + + +def validate_runtime_lora_scope( + *, model_name: str, quantization_name: str, tensor_parallel_shards: int +) -> None: + """Enforce the deliberately narrow scope of the first runtime LoRA implementation.""" + if model_name != "qwen2": + raise ValueError("Runtime LoRA currently supports only the Qwen2 model family.") + if quantization_name != "q0f16": + raise ValueError("Runtime LoRA currently requires `q0f16` model weights.") + if tensor_parallel_shards != 1: + raise ValueError("Runtime LoRA currently requires `tensor_parallel_shards=1`.") + + +__all__ = ["RuntimeLoRAConfig", "validate_runtime_lora_scope"] diff --git a/tests/python/support/test_cli_convert_weight.py b/tests/python/support/test_cli_convert_weight.py index d741087ec2..9bfb276b59 100644 --- a/tests/python/support/test_cli_convert_weight.py +++ b/tests/python/support/test_cli_convert_weight.py @@ -65,5 +65,56 @@ def _fake_convert_weight(**kwargs): ) assert call_args["lora_adapter"] == adapter_dir + assert call_args["lora_mode"] == "merge" assert call_args["source"] == source_index assert call_args["source_format"] == "huggingface-torch" + + +def test_convert_weight_cli_passes_runtime_lora_mode(monkeypatch): + with tempfile.TemporaryDirectory() as tmp_dir: + temp_path = Path(tmp_dir) + config_path = temp_path / "config.json" + source_dir = temp_path / "source" + source_dir.mkdir() + source_index = source_dir / "model.safetensors" + source_index.write_bytes(b"") + adapter_dir = temp_path / "adapter" + adapter_dir.mkdir() + output_dir = temp_path / "output" + config_path.write_text(json.dumps({}), encoding="utf-8") + + monkeypatch.setattr(convert_weight_cli, "detect_config", Path) + monkeypatch.setattr(convert_weight_cli, "detect_device", lambda device: device) + monkeypatch.setattr( + convert_weight_cli, + "detect_weight", + lambda *_args: (source_index, "huggingface-safetensor"), + ) + monkeypatch.setattr(convert_weight_cli, "detect_model_type", lambda *_args: "dummy") + monkeypatch.setattr(convert_weight_cli, "MODELS", {"dummy": object()}) + monkeypatch.setattr(convert_weight_cli, "QUANTIZATION", {"q0f16": object()}) + call_args = {} + monkeypatch.setattr( + convert_weight_cli, + "convert_weight", + lambda **kwargs: call_args.update(kwargs), + ) + + convert_weight_cli.main( + [ + str(config_path), + "--quantization", + "q0f16", + "--model-type", + "dummy", + "--source", + str(source_dir), + "--output", + str(output_dir), + "--lora-adapter", + str(adapter_dir), + "--lora-mode", + "runtime", + ] + ) + assert call_args["lora_mode"] == "runtime" diff --git a/tests/python/support/test_convert_weight_lora_merge.py b/tests/python/support/test_convert_weight_lora_merge.py index e5376466c7..e586684181 100644 --- a/tests/python/support/test_convert_weight_lora_merge.py +++ b/tests/python/support/test_convert_weight_lora_merge.py @@ -106,3 +106,49 @@ def test_convert_weight_with_lora_rejects_awq(): output=temp_path / "output", lora_adapter=adapter_dir, ) + + +def test_convert_weight_runtime_preserves_separate_adapter(monkeypatch): + with tempfile.TemporaryDirectory() as tmp_dir: + temp_path = Path(tmp_dir) + config_path = temp_path / "config.json" + config_path.write_text(json.dumps({}), encoding="utf-8") + adapter_dir = temp_path / "adapter" + adapter_dir.mkdir() + (adapter_dir / "adapter_config.json").write_text( + json.dumps( + { + "bias": "none", + "lora_alpha": 8, + "peft_type": "LORA", + "r": 4, + "target_modules": ["q_proj", "v_proj"], + "task_type": "CAUSAL_LM", + } + ), + encoding="utf-8", + ) + captured = {} + monkeypatch.setattr( + convert_weight_interface, + "_convert_args", + lambda args: captured.update(args=args), + ) + monkeypatch.setattr(convert_weight_interface.ConversionArgs, "display", lambda self: None) + + convert_weight_interface.convert_weight( + config=config_path, + quantization=object(), + model=type("DummyModel", (), {"name": "qwen2"})(), + device=object(), + source=temp_path / "model.safetensors", + source_format="huggingface-safetensor", + output=temp_path / "output", + lora_adapter=adapter_dir, + lora_mode="runtime", + ) + + args = captured["args"] + assert args.lora_mode == "runtime" + assert args.runtime_lora.rank == 4 + assert args.source == temp_path / "model.safetensors" diff --git a/tests/python/support/test_qwen2_runtime_lora.py b/tests/python/support/test_qwen2_runtime_lora.py new file mode 100644 index 0000000000..9978fa18d4 --- /dev/null +++ b/tests/python/support/test_qwen2_runtime_lora.py @@ -0,0 +1,65 @@ +import pytest + +from mlc_llm.model.qwen2 import qwen2_loader +from mlc_llm.model.qwen2.qwen2_model import QWen2Config, QWen2LMHeadModel +from mlc_llm.quantization import QUANTIZATION +from mlc_llm.support.runtime_lora import RuntimeLoRAConfig + +pytestmark = [pytest.mark.unittest] + + +def _config() -> QWen2Config: + return QWen2Config( + hidden_act="silu", + hidden_size=16, + intermediate_size=32, + num_attention_heads=4, + num_hidden_layers=1, + num_key_value_heads=2, + rms_norm_eps=1e-6, + rope_theta=10000, + vocab_size=32, + context_window_size=64, + prefill_chunk_size=16, + runtime_lora=RuntimeLoRAConfig( + rank=4, + alpha=8.0, + target_modules=("down_proj", "gate_proj", "q_proj", "v_proj"), + ), + ) + + +def test_qwen2_exports_separate_runtime_lora_parameters(): + config = _config() + model = QWen2LMHeadModel(config) + model.to("float16") + _, exported_params, _ = model.export_tvm(spec=model.get_default_spec(), allow_extern=True) + runtime_params = { + name: param for name, param in exported_params if param.attrs.get("runtime_lora_param") + } + + assert len(runtime_params) == 8 + assert ( + runtime_params["model.layers.0.self_attn.q_proj_lora.lora_a.weight"].attrs[ + "peft_source_suffix" + ] + == "model.layers.0.self_attn.q_proj.lora_A.weight" + ) + assert ( + runtime_params["model.layers.0.mlp.down_proj_lora.lora_b.weight"].attrs[ + "peft_source_suffix" + ] + == "model.layers.0.mlp.down_proj.lora_B.weight" + ) + assert ( + runtime_params["model.layers.0.mlp.down_proj_lora.lora_b.weight"].attrs[ + "runtime_lora_scale" + ] + == 2.0 + ) + + +def test_qwen2_base_loader_does_not_claim_adapter_parameters(): + mapping = qwen2_loader.huggingface(_config(), QUANTIZATION["q0f16"]) + assert not any("_lora." in name for name in mapping.param_map) + assert "model.layers.0.self_attn.c_attn.weight" in mapping.param_map diff --git a/tests/python/support/test_runtime_lora.py b/tests/python/support/test_runtime_lora.py new file mode 100644 index 0000000000..a2705d6e41 --- /dev/null +++ b/tests/python/support/test_runtime_lora.py @@ -0,0 +1,158 @@ +import json +import math +import tempfile +from pathlib import Path + +import numpy as np +import pytest + +from mlc_llm.loader.runtime_lora import ( + make_runtime_lora_mapping, + resolve_runtime_lora_weight, +) +from mlc_llm.support.runtime_lora import RuntimeLoRAConfig, validate_runtime_lora_scope + +pytestmark = [pytest.mark.unittest] + + +def _write_adapter_config(adapter_dir: Path, **overrides) -> None: + config = { + "bias": "none", + "fan_in_fan_out": False, + "lora_alpha": 16, + "peft_type": "LORA", + "r": 8, + "target_modules": ["q_proj", "v_proj"], + "task_type": "CAUSAL_LM", + "use_dora": False, + "use_rslora": False, + } + config.update(overrides) + (adapter_dir / "adapter_config.json").write_text(json.dumps(config), encoding="utf-8") + + +def _write_safetensor_header(path: Path, names) -> None: + header = {name: {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]} for name in names} + header_bytes = json.dumps(header).encode("utf-8") + path.write_bytes(len(header_bytes).to_bytes(8, "little") + header_bytes) + + +def test_runtime_lora_config_standard_and_rslora_scaling(): + with tempfile.TemporaryDirectory() as tmp_dir: + adapter_dir = Path(tmp_dir) + _write_adapter_config(adapter_dir) + config = RuntimeLoRAConfig.from_peft_directory(adapter_dir) + assert config.rank == 8 + assert config.target_modules == ("q_proj", "v_proj") + assert config.scaling == 2.0 + + _write_adapter_config(adapter_dir, use_rslora=True) + config = RuntimeLoRAConfig.from_peft_directory(adapter_dir) + assert config.scaling == pytest.approx(16 / math.sqrt(8)) + assert RuntimeLoRAConfig.from_dict(config.asdict()) == config + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"target_modules": ["lm_head"]}, "unsupported"), + ({"rank_pattern": {"q_proj": 4}}, "rank_pattern"), + ({"alora_invocation_tokens": [1, 2]}, "alora_invocation_tokens"), + ({"arrow_config": {"top_k": 1}}, "arrow_config"), + ({"exclude_modules": ["k_proj"]}, "exclude_modules"), + ({"lora_bias": True}, "lora_bias"), + ({"use_qalora": True}, "use_qalora"), + ({"use_dora": True}, "DoRA"), + ({"bias": "all"}, "bias"), + ], +) +def test_runtime_lora_config_rejects_unsupported_peft_features(overrides, message): + with tempfile.TemporaryDirectory() as tmp_dir: + adapter_dir = Path(tmp_dir) + _write_adapter_config(adapter_dir, **overrides) + with pytest.raises(ValueError, match=message): + RuntimeLoRAConfig.from_peft_directory(adapter_dir) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"lora_alpha": float("inf")}, "finite"), + ({"target_modules": ["q_proj", 1]}, "string"), + ({"use_rslora": "false"}, "use_rslora"), + ], +) +def test_runtime_lora_config_rejects_malformed_values(overrides, message): + with tempfile.TemporaryDirectory() as tmp_dir: + adapter_dir = Path(tmp_dir) + _write_adapter_config(adapter_dir, **overrides) + with pytest.raises(ValueError, match=message): + RuntimeLoRAConfig.from_peft_directory(adapter_dir) + + +def test_runtime_lora_scope_is_deliberately_narrow(): + validate_runtime_lora_scope( + model_name="qwen2", quantization_name="q0f16", tensor_parallel_shards=1 + ) + with pytest.raises(ValueError, match="Qwen2"): + validate_runtime_lora_scope( + model_name="llama", quantization_name="q0f16", tensor_parallel_shards=1 + ) + with pytest.raises(ValueError, match="q0f16"): + validate_runtime_lora_scope( + model_name="qwen2", quantization_name="q4f16_1", tensor_parallel_shards=1 + ) + with pytest.raises(ValueError, match="tensor_parallel_shards"): + validate_runtime_lora_scope( + model_name="qwen2", quantization_name="q0f16", tensor_parallel_shards=2 + ) + + +def test_runtime_lora_mapping_matches_peft_prefix_and_adapter_name(): + class _FakeParam: + dtype = "float16" + + def __init__(self, suffix, scale=1.0): + self.attrs = {"peft_source_suffix": suffix} + if scale != 1.0: + self.attrs["runtime_lora_scale"] = scale + + with tempfile.TemporaryDirectory() as tmp_dir: + adapter_dir = Path(tmp_dir) + adapter_weight = adapter_dir / "adapter_model.safetensors" + a_name = "base_model.model.model.layers.0.self_attn.q_proj.lora_A.default.weight" + b_name = "base_model.model.model.layers.0.self_attn.q_proj.lora_B.default.weight" + _write_safetensor_header(adapter_weight, [a_name, b_name]) + + named_parameters = { + "model.layers.0.self_attn.q_proj_lora.lora_a.weight": _FakeParam( + "model.layers.0.self_attn.q_proj.lora_A.weight" + ), + "model.layers.0.self_attn.q_proj_lora.lora_b.weight": _FakeParam( + "model.layers.0.self_attn.q_proj.lora_B.weight", scale=0.1 + ), + } + mapping = make_runtime_lora_mapping(named_parameters, adapter_weight) + assert resolve_runtime_lora_weight(adapter_dir) == adapter_weight + assert mapping.param_map["model.layers.0.self_attn.q_proj_lora.lora_a.weight"] == [a_name] + assert mapping.param_map["model.layers.0.self_attn.q_proj_lora.lora_b.weight"] == [b_name] + weight = np.array([-0.13210486, -0.53731406, 1.9878846, -0.29650125]).reshape(2, 2) + mapped_a = mapping.map_func["model.layers.0.self_attn.q_proj_lora.lora_a.weight"](weight) + mapped_b = mapping.map_func["model.layers.0.self_attn.q_proj_lora.lora_b.weight"](weight) + np.testing.assert_array_equal(mapped_a, weight.astype("float16")) + np.testing.assert_array_equal(mapped_b, (weight * 0.1).astype("float16")) + assert not np.array_equal(mapped_b, weight.astype("float16") * np.float16(0.1)) + + +def test_runtime_lora_math_matches_merged_weight(): + rng = np.random.default_rng(0) + x = rng.normal(size=(2, 3, 8)).astype("float32") + weight = rng.normal(size=(6, 8)).astype("float32") + lora_a = rng.normal(size=(4, 8)).astype("float32") + lora_b = rng.normal(size=(6, 4)).astype("float32") + scale = 3.0 + + scaled_lora_b = lora_b * scale + runtime_output = x @ weight.T + (x @ lora_a.T) @ scaled_lora_b.T + merged_output = x @ (weight + scale * (lora_b @ lora_a)).T + np.testing.assert_allclose(runtime_output, merged_output, rtol=1e-5, atol=1e-5) From b8e7cf93d175040b2361a1ba28c4703cb61d0304 Mon Sep 17 00:00:00 2001 From: TR-3B <144127816+MagellaX@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:35:01 +0530 Subject: [PATCH 2/2] Update python/mlc_llm/model/qwen2/qwen2_model.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- python/mlc_llm/model/qwen2/qwen2_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mlc_llm/model/qwen2/qwen2_model.py b/python/mlc_llm/model/qwen2/qwen2_model.py index 3e620713b4..59540c6b41 100644 --- a/python/mlc_llm/model/qwen2/qwen2_model.py +++ b/python/mlc_llm/model/qwen2/qwen2_model.py @@ -175,7 +175,7 @@ def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: k = k + self.k_proj_lora(hidden_states) if self.v_proj_lora is not None: v = v + self.v_proj_lora(hidden_states) - qkv = op.concat([q, k, v], dim=-1) + qkv = op.concat([q, k, v], axis=-1) qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) output = op.reshape( paged_kv_cache.attention_with_fused_qkv(