diff --git a/tests/lora/test_lora_weights.py b/tests/lora/test_lora_weights.py new file mode 100644 index 0000000000..ff13214793 --- /dev/null +++ b/tests/lora/test_lora_weights.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights + +RANK, ALPHA = 16, 32 +# PEFTHelper stores alpha/sqrt(r) for use_rslora adapters. +RSLORA_SCALING = ALPHA / RANK**0.5 + + +@pytest.fixture(autouse=True) +def should_do_global_cleanup_after_test(): + # Pure-python weight packing; skip GPU/ray cleanup. + return False + + +def moe_loras(scaling: float, stacked: bool = False) -> list[LoRALayerWeights]: + a_shape = (3, RANK, 4) if stacked else (RANK, 4) + b_shape = (3, 6, RANK) if stacked else (6, RANK) + return [ + LoRALayerWeights( + module_name=f"experts.w{i}", + rank=RANK, + lora_alpha=ALPHA, + lora_a=torch.randn(*a_shape), + lora_b=torch.randn(*b_shape), + scaling=scaling, + ) + for i in (1, 2, 3) + ] + + +@pytest.mark.parametrize( + "packer,stacked", + [ + (PackedLoRALayerWeights.pack_moe, False), + (PackedLoRALayerWeights.pack_moe_stacked, True), + ], +) +def test_moe_packing_uses_stored_scaling(packer, stacked: bool): + packed = packer(moe_loras(RSLORA_SCALING, stacked), "experts") + assert packed.scaling == pytest.approx([RSLORA_SCALING] * 3) + + +def test_non_gated_moe_keeps_w3_unscaled(): + packed = PackedLoRALayerWeights.pack_moe( + moe_loras(RSLORA_SCALING), "experts", is_non_gated_moe=True + ) + assert packed.scaling[:2] == pytest.approx([RSLORA_SCALING] * 2) + assert packed.scaling[2] == 1.0 diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index 38a5875525..bb241db77e 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -201,8 +201,9 @@ 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. - scaling = lora_alpha / rank + # 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 if is_non_gated_moe: @@ -250,7 +251,7 @@ 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 = lora_alpha / rank + scaling = w1_lora.scaling return cls( module_name, rank,