[Bugfix][LoRA] Apply PEFT rank_pattern / alpha_pattern scaling per module (#59801)

Signed-off-by: suraj <surajsharan85@gmail.com>
Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
Co-authored-by: linitra24 <renshuang.zhou@daocloud.io>

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)

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},
+        )

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:

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":

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":