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