[Bugfix][LoRA] Use stored rsLoRA scaling factor in MoE expert packing (#55548)
Signed-off-by: Kushal <72650064+kushaldabbe@users.noreply.github.com>
Kushal committed
Sep 11, 2026 at 12:52 UTC
b5d4186300fda544ec5ec51113d569159c5c52bd
2 files changed
+56
-3
tests/lora/test_lora_weights.py
new
+52
@@ -0,0 +1,52 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
7
+
8
+RANK, ALPHA = 16, 32
9
+# PEFTHelper stores alpha/sqrt(r) for use_rslora adapters.
10
+RSLORA_SCALING = ALPHA / RANK**0.5
11
+
12
+
13
+@pytest.fixture(autouse=True)
14
+def should_do_global_cleanup_after_test():
15
+ # Pure-python weight packing; skip GPU/ray cleanup.
16
+ return False
17
+
18
+
19
+def moe_loras(scaling: float, stacked: bool = False) -> list[LoRALayerWeights]:
20
+ a_shape = (3, RANK, 4) if stacked else (RANK, 4)
21
+ b_shape = (3, 6, RANK) if stacked else (6, RANK)
22
+ return [
23
+ LoRALayerWeights(
24
+ module_name=f"experts.w{i}",
25
+ rank=RANK,
26
+ lora_alpha=ALPHA,
27
+ lora_a=torch.randn(*a_shape),
28
+ lora_b=torch.randn(*b_shape),
29
+ scaling=scaling,
30
+ )
31
+ for i in (1, 2, 3)
32
+ ]
33
+
34
+
35
+@pytest.mark.parametrize(
36
+ "packer,stacked",
37
+ [
38
+ (PackedLoRALayerWeights.pack_moe, False),
39
+ (PackedLoRALayerWeights.pack_moe_stacked, True),
40
+ ],
41
+)
42
+def test_moe_packing_uses_stored_scaling(packer, stacked: bool):
43
+ packed = packer(moe_loras(RSLORA_SCALING, stacked), "experts")
44
+ assert packed.scaling == pytest.approx([RSLORA_SCALING] * 3)
45
+
46
+
47
+def test_non_gated_moe_keeps_w3_unscaled():
48
+ packed = PackedLoRALayerWeights.pack_moe(
49
+ moe_loras(RSLORA_SCALING), "experts", is_non_gated_moe=True
50
+ )
51
+ assert packed.scaling[:2] == pytest.approx([RSLORA_SCALING] * 2)
52
+ assert packed.scaling[2] == 1.0
vllm/lora/lora_weights.py
+4
-3
@@ -201,8 +201,9 @@ class PackedLoRALayerWeights(LoRALayerWeights):
201
w1_lora_b = torch.stack(w1_lora_b_lst, dim=0) # (num_experts,output_size,rank)
202
w2_lora_b = torch.stack(w2_lora_b_lst, dim=0)
203
204
- # All w1, w2, w3 have the same scaling factor.
205
- scaling = lora_alpha / rank
204
+ # All w1, w2, w3 have the same scaling factor. Use the per-adapter
205
+ # scaling (e.g. alpha/sqrt(r) for rsLoRA) instead of alpha/rank.
206
+ scaling = first_lora.scaling
207
last_scaling = scaling
208
209
if is_non_gated_moe:
@@ -250,7 +251,7 @@ class PackedLoRALayerWeights(LoRALayerWeights):
251
assert w1_lora is not None and w2_lora is not None and w3_lora is not None
252
rank = w1_lora.rank
253
lora_alpha = w1_lora.lora_alpha
253
- scaling = lora_alpha / rank
254
+ scaling = w1_lora.scaling
255
return cls(
256
module_name,
257
rank,