From 4c36d21280824d1a9c521106932408243078de0d Mon Sep 17 00:00:00 2001 From: Suraj Sharan Date: Mon, 12 Oct 2026 17:49:30 +0400 Subject: [PATCH] [Bugfix][LoRA] Apply PEFT rank_pattern / alpha_pattern scaling per module (#59801) Signed-off-by: suraj Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Co-authored-by: linitra24 --- tests/lora/test_lora_weights.py | 57 +++++++++++++++++ tests/lora/test_peft_helper.py | 105 ++++++++++++++++++++++++++++++++ vllm/lora/lora_model.py | 4 +- vllm/lora/lora_weights.py | 38 +++++++++--- vllm/lora/peft_helper.py | 60 ++++++++++++++++-- 5 files changed, 249 insertions(+), 15 deletions(-) diff --git a/tests/lora/test_lora_weights.py b/tests/lora/test_lora_weights.py index ff13214793..7230126e40 100644 --- a/tests/lora/test_lora_weights.py +++ b/tests/lora/test_lora_weights.py @@ -4,6 +4,7 @@ import pytest import torch from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights +from vllm.lora.peft_helper import PEFTHelper RANK, ALPHA = 16, 32 # PEFTHelper stores alpha/sqrt(r) for use_rslora adapters. @@ -50,3 +51,59 @@ def test_non_gated_moe_keeps_w3_unscaled(): ) assert packed.scaling[:2] == pytest.approx([RSLORA_SCALING] * 2) assert packed.scaling[2] == 1.0 + + +@pytest.mark.parametrize( + "packer,stacked", + [ + (PackedLoRALayerWeights.pack_moe, False), + (PackedLoRALayerWeights.pack_moe_stacked, True), + ], +) +def test_moe_packing_keeps_per_projection_scaling(packer, stacked: bool): + # e.g. rank_pattern / alpha_pattern giving gate, down and up different scalings + loras = moe_loras(1.0, stacked) + for lora, scaling in zip(loras, (2.0, 4.0, 8.0)): + lora.scaling = scaling + packed = packer(loras, "experts") + assert packed.scaling == pytest.approx([2.0, 4.0, 8.0]) + + +def test_moe_packing_keeps_per_expert_scaling(): + """Experts of one projection can have different scaling, e.g. from + alpha_pattern.""" + loras = moe_loras(RSLORA_SCALING) + moe_loras(RSLORA_SCALING) + loras[4].scaling = 2 * RSLORA_SCALING # w2 of expert 1 + expected = [lora.lora_b * lora.scaling for lora in loras] + packed = PackedLoRALayerWeights.pack_moe(loras, "experts").optimize() + for i, lora_b in enumerate(packed.lora_b): + torch.testing.assert_close(lora_b, torch.stack(expected[i::3])) + + +def test_moe_packing_rejects_per_expert_rank(): + """Experts of one projection must share a rank to be stacked.""" + loras = moe_loras(1.0) + moe_loras(1.0) + loras[4].lora_a = torch.randn(RANK // 2, 4) + loras[4].lora_b = torch.randn(6, RANK // 2) + with pytest.raises(ValueError, match="different LoRA ranks"): + PackedLoRALayerWeights.pack_moe(loras, "experts") + + +def test_from_config_uses_checkpoint_module_name(): + peft_helper = PEFTHelper( + r=RANK, + lora_alpha=ALPHA, + target_modules=["q_proj"], + rank_pattern={r"^model\.layers\.0\.self_attn\.q_proj": 4}, + ) + # the pattern refers to the checkpoint name, not the name after weights mapping + lora = LoRALayerWeights.from_config( + "language_model.model.layers.0.self_attn.q_proj", + peft_helper, + "model.layers.0.self_attn.q_proj", + ) + assert lora.rank == 4 + assert lora.scaling == pytest.approx(ALPHA / 4) + lora = LoRALayerWeights.from_config("model.layers.1.self_attn.q_proj", peft_helper) + assert lora.rank == RANK + assert lora.scaling == pytest.approx(ALPHA / RANK) diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index bb1246144b..2fcf5ff962 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -22,6 +22,11 @@ ERROR_CASES = [ {"modules_to_save": ["lm_head"]}, "Unsupported modules_to_save", ), + ( + "test_rank_pattern", + {"r": 8, "rank_pattern": {"q_proj": 32}}, + "LoRA rank 32 is greater than max_lora_rank", + ), ("test_rank_zero", {"r": 0}, "must be a positive integer"), ("test_rank_negative", {"r": -8}, "must be a positive integer"), ("test_lora_bias", {"lora_bias": True}, "does not support LoRA bias"), @@ -153,3 +158,103 @@ def test_peft_helper_init_lora_weights_supported(init_lora_weights): "use_bdlora": None, } ).validate_legal(lora_config) + + +@pytest.mark.parametrize( + "config,module_name,expected", + [ + # rank_pattern only: alpha / r_m + ({"rank_pattern": {"q_proj": 4}}, "model.layers.0.self_attn.q_proj", (4, 8.0)), + ({"rank_pattern": {"q_proj": 4}}, "model.layers.0.self_attn.v_proj", (16, 2.0)), + # alpha_pattern only: alpha_m / r + ( + {"alpha_pattern": {"v_proj": 64}}, + "model.layers.0.self_attn.v_proj", + (16, 4.0), + ), + # keys only match at a "." boundary, like in PEFT + ({"rank_pattern": {"proj": 4}}, "model.layers.0.self_attn.q_proj", (16, 2.0)), + # fused MoE experts from PEFT `target_parameters`: gate_up_proj is stored + # as `experts.base_layer`, down_proj as `experts` + ( + {"rank_pattern": {"experts.gate_up_proj": 2}}, + "model.layers.0.mlp.experts.base_layer", + (2, 16.0), + ), + ( + {"rank_pattern": {"experts.down_proj": 2}}, + "model.layers.0.mlp.experts", + (2, 16.0), + ), + ( + {"rank_pattern": {"experts.gate_up_proj": 2}}, + "model.layers.0.mlp.experts", + (16, 2.0), + ), + ], +) +def test_peft_helper_rank_and_alpha_pattern(config, module_name, expected): + peft_helper = PEFTHelper( + r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"], **config + ) + rank, scaling = peft_helper.get_rank_and_scaling(module_name) + assert rank == expected[0] + assert scaling == pytest.approx(expected[1]) + + +def test_peft_helper_pattern_rslora(): + peft_helper = PEFTHelper( + r=16, + lora_alpha=32, + target_modules=["q_proj"], + use_rslora=True, + rank_pattern={"q_proj": 4}, + ) + assert peft_helper.get_rank_and_scaling("model.q_proj") == (4, 32 / math.sqrt(4)) + + +def test_peft_helper_dynamic_rank_conversion(): + """PEFT's `save_as_lora` with a dynamic rank writes r=1, lora_alpha=1 and the + same per-module values in rank_pattern and alpha_pattern, so every module + keeps a scaling of 1.""" + peft_helper = PEFTHelper( + r=1, + lora_alpha=1, + target_modules=["q_proj", "v_proj"], + rank_pattern={"layers.0.self_attn.q_proj": 32, "layers.0.self_attn.v_proj": 3}, + alpha_pattern={"layers.0.self_attn.q_proj": 32, "layers.0.self_attn.v_proj": 3}, + ) + assert peft_helper.get_rank_and_scaling("model.layers.0.self_attn.q_proj") == ( + 32, + 1.0, + ) + assert peft_helper.get_rank_and_scaling("model.layers.0.self_attn.v_proj") == ( + 3, + 1.0, + ) + with pytest.raises(ValueError, match="LoRA rank 32 is greater than max_lora_rank"): + peft_helper.validate_legal(LoRAConfig(max_lora_rank=16)) + + +def test_peft_helper_null_patterns(): + peft_helper = PEFTHelper.from_dict( + { + "r": 8, + "lora_alpha": 16, + "target_modules": ["q_proj"], + "rank_pattern": None, + "alpha_pattern": None, + } + ) + assert peft_helper.get_rank_and_scaling("model.q_proj") == (8, 2.0) + + +@pytest.mark.parametrize("bad_rank", [0, -4, 2.5]) +def test_peft_helper_invalid_rank_pattern(bad_rank): + with pytest.raises(ValueError, match="rank_pattern"): + PEFTHelper( + r=8, + lora_alpha=16, + target_modules=["q_proj"], + rank_pattern={"q_proj": bad_rank}, + ) diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index 1964453324..f4520b12c7 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -150,8 +150,10 @@ class LoRAModel: continue if module_name not in loras: + # rank/alpha patterns refer to the names before weights mapping + peft_module_name, _ = parse_fine_tuned_lora_name(tensor_name) loras[module_name] = LoRALayerWeights.from_config( - module_name, peft_helper + module_name, peft_helper, peft_module_name ) if is_lora_a: diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index 13a692a623..9ea7c550d7 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -59,15 +59,22 @@ class LoRALayerWeights: cls, module_name: str, peft_helper: PEFTHelper, + peft_module_name: str | None = None, ) -> "LoRALayerWeights": + """`peft_module_name` is the module name in the adapter checkpoint, used + to resolve `rank_pattern` / `alpha_pattern`. Defaults to `module_name`. + """ + rank, scaling = peft_helper.get_rank_and_scaling( + peft_module_name or module_name + ) # lora_a and lora_b are set to None for config-based construction return cls( module_name, - peft_helper.r, + rank, peft_helper.lora_alpha, None, None, - peft_helper.vllm_lora_scaling_factor, + scaling, ) @classmethod @@ -168,6 +175,18 @@ class PackedLoRALayerWeights(LoRALayerWeights): rank = first_lora.rank lora_alpha = first_lora.lora_alpha assert len(loras) % 3 == 0 + for proj_loras in (loras[0::3], loras[1::3], loras[2::3]): + experts = [lora for lora in proj_loras if lora is not None] + if len({lora.lora_a.shape[0] for lora in experts}) > 1: + raise ValueError( + f"The experts of {module_name} have different LoRA ranks for " + "the same projection, which is not supported." + ) + # The packed LoRA has one scaling per projection, so fold it into + # lora_b where the experts differ (e.g. from alpha_pattern). + if len({lora.scaling for lora in experts}) > 1: + for lora in experts: + lora.optimize() w1_lora_a_lst = [] w2_lora_a_lst = [] w3_lora_a_lst = [] @@ -201,10 +220,12 @@ class PackedLoRALayerWeights(LoRALayerWeights): w1_lora_b = torch.stack(w1_lora_b_lst, dim=0) # (num_experts,output_size,rank) w2_lora_b = torch.stack(w2_lora_b_lst, dim=0) - # All w1, w2, w3 have the same scaling factor. Use the per-adapter - # scaling (e.g. alpha/sqrt(r) for rsLoRA) instead of alpha/rank. - scaling = first_lora.scaling - last_scaling = scaling + # Use each projection's own scaling (e.g. alpha/sqrt(r) for rsLoRA, or + # per-module values from rank_pattern / alpha_pattern) instead of + # alpha/rank. + w1_scaling = loras[0].scaling # type: ignore[union-attr] + w2_scaling = loras[1].scaling # type: ignore[union-attr] + last_scaling = loras[2].scaling if loras[2] is not None else w1_scaling if is_non_gated_moe: # For non-gated MoE, reuse w1 tensors for w3 to avoid memory waste @@ -224,7 +245,7 @@ class PackedLoRALayerWeights(LoRALayerWeights): [lora_alpha, lora_alpha, lora_alpha], [w1_lora_a, w2_lora_a, w3_lora_a], [w1_lora_b, w2_lora_b, w3_lora_b], - scaling=[scaling, scaling, last_scaling], + scaling=[w1_scaling, w2_scaling, last_scaling], ) return obj @@ -251,14 +272,13 @@ class PackedLoRALayerWeights(LoRALayerWeights): assert w1_lora is not None and w2_lora is not None and w3_lora is not None rank = w1_lora.rank lora_alpha = w1_lora.lora_alpha - scaling = w1_lora.scaling return cls( module_name, rank, [lora_alpha, lora_alpha, lora_alpha], [w1_lora.lora_a, w2_lora.lora_a, w3_lora.lora_a], [w1_lora.lora_b, w2_lora.lora_b, w3_lora.lora_b], - scaling=[scaling, scaling, scaling], + scaling=[w1_lora.scaling, w2_lora.scaling, w3_lora.scaling], ) def optimize(self) -> "PackedLoRALayerWeights": diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index de05dea1cf..721b7a039f 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -9,6 +9,8 @@ import os from dataclasses import MISSING, dataclass, field, fields from typing import Literal +import regex as re + from vllm.config.lora import LoRAConfig from vllm.logger import init_logger @@ -41,6 +43,9 @@ class PEFTHelper: layer_replication: list[list[int]] | None = field(default=None) use_bdlora: dict | None = field(default=None) use_qalora: bool = field(default=False) + # Per-module overrides of `r` and `lora_alpha`, keyed by module name pattern + rank_pattern: dict[str, int] = field(default_factory=dict) + alpha_pattern: dict[str, float] = field(default_factory=dict) # Extra vllm field, start with 'vllm_' to avoid conflict vllm_lora_scaling_factor: float = field(default=1.0) vllm_max_position_embeddings: int | None = field(default=False) @@ -97,11 +102,55 @@ class PEFTHelper: def __post_init__(self): if self.r <= 0: raise ValueError(f"LoRA rank `r` must be a positive integer, got {self.r}.") + self.rank_pattern = self.rank_pattern or {} + self.alpha_pattern = self.alpha_pattern or {} + for key, rank in self.rank_pattern.items(): + if not isinstance(rank, int) or rank <= 0: + raise ValueError( + f"LoRA rank for `{key}` in `rank_pattern` must be a positive " + f"integer, got {rank}." + ) if self.use_rslora: logger.info_once("Loading LoRA weights trained with rsLoRA.") - self.vllm_lora_scaling_factor = self.lora_alpha / math.sqrt(self.r) - else: - self.vllm_lora_scaling_factor = self.lora_alpha / self.r + self.vllm_lora_scaling_factor = self._scaling(self.r, self.lora_alpha) + + def _scaling(self, rank: int, alpha: float) -> float: + if self.use_rslora: + return alpha / math.sqrt(rank) + return alpha / rank + + @staticmethod + def _match_pattern(pattern: dict, name: str): + # Same rule as PEFT's `get_pattern_key`: a key matches if it is a suffix + # of the module name starting at a "." boundary; the first match wins. + for key, value in pattern.items(): + if re.match(rf"(.*\.)?({key})$", name): + return value + return None + + def get_rank_and_scaling(self, module_name: str) -> tuple[int, float]: + """Return the rank and scaling PEFT uses for `module_name`, taking + `rank_pattern` and `alpha_pattern` into account. + + `module_name` is the name in the adapter checkpoint, without the + `base_model.model.` prefix and before any weights mapping. + """ + if not self.rank_pattern and not self.alpha_pattern: + return self.r, self.vllm_lora_scaling_factor + # PEFT stores fused MoE expert LoRAs (`target_parameters`) as + # `experts.base_layer` (gate_up_proj) and `experts` (down_proj), but + # matches the patterns against the parameter names. + if module_name.endswith(".experts.base_layer"): + module_name = module_name.removesuffix(".base_layer") + ".gate_up_proj" + elif module_name.endswith(".experts"): + module_name += ".down_proj" + rank = self._match_pattern(self.rank_pattern, module_name) + if rank is None: + rank = self.r + alpha = self._match_pattern(self.alpha_pattern, module_name) + if alpha is None: + alpha = self.lora_alpha + return rank, self._scaling(rank, alpha) @classmethod def from_dict(cls, config_dict: dict) -> "PEFTHelper": @@ -141,9 +190,10 @@ class PEFTHelper: constraints and requirements. """ error_msg = self._validate_features() - if self.r > lora_config.max_lora_rank: + max_rank = max([self.r, *self.rank_pattern.values()]) + if max_rank > lora_config.max_lora_rank: error_msg.append( - f"LoRA rank {self.r} is greater than max_lora_rank" + f"LoRA rank {max_rank} is greater than max_lora_rank" f" {lora_config.max_lora_rank}." ) if self.bias != "none":