| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """End-to-end correctness test for 2D MoE LoRA expert-parallel |
| 4 | load-time slicing |
| 5 | """ |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec |
| 11 | from vllm.lora.peft_helper import PEFTHelper |
| 12 | |
| 13 | NUM_LAYERS = 48 |
| 14 | GLOBAL_NUM_EXPERTS = 128 |
| 15 | LOCAL_NUM_EXPERTS = 64 # ep_size = 2 |
| 16 | EXPERT_PROJECTIONS = ("down_proj", "gate_proj", "up_proj") |
| 17 | NON_EXPERT_MODULES = ("q_proj", "k_proj", "v_proj", "o_proj", "gate") |
| 18 | |
| 19 | |
| 20 | def _expected_lora_modules() -> set[str]: |
| 21 | """Replicate the set ``WorkerLoRAManager._load_adapter`` would build |
| 22 | from this model's ``packed_modules_mapping``.""" |
| 23 | expected: set[str] = set(NON_EXPERT_MODULES) |
| 24 | for expert in range(GLOBAL_NUM_EXPERTS): |
| 25 | for proj in EXPERT_PROJECTIONS: |
| 26 | expected.add(f"experts.{expert}.{proj}") |
| 27 | return expected |
| 28 | |
| 29 | |
| 30 | def _load(lora_dir, peft_helper, *, moe_ep_spec, lora_id): |
| 31 | return LoRAModel.from_local_checkpoint( |
| 32 | lora_dir, |
| 33 | _expected_lora_modules(), |
| 34 | peft_helper=peft_helper, |
| 35 | lora_model_id=lora_id, |
| 36 | device="cpu", |
| 37 | moe_ep_spec=moe_ep_spec, |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | @pytest.mark.parametrize("ep_rank", [0, 1]) |
| 42 | def test_moe_lora_ep2_real_qwen3moe(qwen3moe_lora_files, ep_rank): |
| 43 | """ep_size=2 against the real Qwen3-MoE adapter: each rank's loaded |
| 44 | LoRA has the right size, the right expert membership, and the |
| 45 | right tensor values.""" |
| 46 | peft_helper = PEFTHelper.from_local_dir( |
| 47 | qwen3moe_lora_files, max_position_embeddings=4096 |
| 48 | ) |
| 49 | |
| 50 | # Baseline: no spec → loads every expert × projection × layer plus |
| 51 | # all non-expert LoRA modules. |
| 52 | ground_truth = _load(qwen3moe_lora_files, peft_helper, moe_ep_spec=None, lora_id=1) |
| 53 | expected_baseline = ( |
| 54 | GLOBAL_NUM_EXPERTS * len(EXPERT_PROJECTIONS) * NUM_LAYERS |
| 55 | + len(NON_EXPERT_MODULES) * NUM_LAYERS |
| 56 | ) |
| 57 | assert len(ground_truth.loras) == expected_baseline |
| 58 | |
| 59 | # Sliced load: only this rank's experts; non-expert LoRA is untouched. |
| 60 | spec = MoEEPLoadSpec( |
| 61 | ep_rank=ep_rank, |
| 62 | local_num_experts=LOCAL_NUM_EXPERTS, |
| 63 | global_num_experts=GLOBAL_NUM_EXPERTS, |
| 64 | ) |
| 65 | sliced = _load( |
| 66 | qwen3moe_lora_files, |
| 67 | peft_helper, |
| 68 | moe_ep_spec=spec, |
| 69 | lora_id=100 + ep_rank, |
| 70 | ) |
| 71 | |
| 72 | expected_sliced = ( |
| 73 | LOCAL_NUM_EXPERTS * len(EXPERT_PROJECTIONS) * NUM_LAYERS |
| 74 | + len(NON_EXPERT_MODULES) * NUM_LAYERS |
| 75 | ) |
| 76 | assert len(sliced.loras) == expected_sliced |
| 77 | |
| 78 | expert_start = ep_rank * LOCAL_NUM_EXPERTS |
| 79 | expert_end = expert_start + LOCAL_NUM_EXPERTS |
| 80 | |
| 81 | for name, lora in sliced.loras.items(): |
| 82 | gt = ground_truth.loras[name] |
| 83 | torch.testing.assert_close(lora.lora_a, gt.lora_a) |
| 84 | torch.testing.assert_close(lora.lora_b, gt.lora_b) |
| 85 | if ".experts." in name: |
| 86 | expert_idx = int(name.split(".experts.")[-1].split(".")[0]) |
| 87 | assert expert_start <= expert_idx < expert_end, ( |
| 88 | f"non-local expert {expert_idx} leaked: {name}" |
| 89 | ) |