@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 156 lines 4.67 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 import vllm
7 import vllm.config
8 from vllm.assets.image import ImageAsset
9 from vllm.lora.request import LoRARequest
10
11 from ..utils import multi_gpu_test
12
13 MODEL_PATH = "Qwen/Qwen3.6-35B-A3B"
14
15 LORA_2D_ID = 1
16 LORA_3D_ID = 2
17
18 PROMPT_TEMPLATE = """<|im_start|>user
19 <|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>
20 <|im_start|>assistant
21 <think>
22
23 </think>
24 """
25
26 # Visual captioning prompts: each image will be paired with one LoRA in the
27 # mixed-batch case so we can check per-prompt routing.
28 VL_TEST_IMAGES = [
29 ImageAsset("stop_sign"),
30 ImageAsset("cherry_blossom"),
31 ]
32
33
34 def _build_prompts() -> list[dict]:
35 return [
36 {
37 "prompt": PROMPT_TEMPLATE,
38 "multi_modal_data": {"image": asset.pil_image},
39 }
40 for asset in VL_TEST_IMAGES
41 ]
42
43
44 def _generate(llm: vllm.LLM, lora_request) -> list[str]:
45 outputs = llm.generate(
46 _build_prompts(),
47 vllm.SamplingParams(temperature=0, max_tokens=128),
48 lora_request=lora_request,
49 )
50 return [out.outputs[0].text.strip() for out in outputs]
51
52
53 def _run_mixed_2d_3d_lora_test(
54 lora_2d_files: str,
55 lora_3d_files: str,
56 tensor_parallel_size: int,
57 fully_sharded_loras: bool,
58 ) -> None:
59 llm = vllm.LLM(
60 model=MODEL_PATH,
61 max_model_len=4096,
62 enable_lora=True,
63 enable_mixed_moe_lora_format=True,
64 max_loras=2,
65 max_lora_rank=8,
66 max_num_seqs=4,
67 enforce_eager=True,
68 tensor_parallel_size=tensor_parallel_size,
69 enable_expert_parallel=not fully_sharded_loras,
70 fully_sharded_loras=fully_sharded_loras,
71 trust_remote_code=True,
72 enable_tower_connector_lora=True,
73 mm_processor_cache_gb=0,
74 limit_mm_per_prompt={"image": 1},
75 compilation_config=vllm.config.CompilationConfig(
76 cudagraph_specialize_lora=False,
77 ),
78 )
79
80 lora_2d = LoRARequest(
81 "lora_2d",
82 LORA_2D_ID,
83 lora_2d_files,
84 is_3d_lora_weight=False,
85 )
86 lora_3d = LoRARequest(
87 "lora_3d",
88 LORA_3D_ID,
89 lora_3d_files,
90 is_3d_lora_weight=True,
91 )
92
93 # Reference: each adapter alone over both prompts.
94 outputs_2d_alone = _generate(llm, lora_2d)
95 outputs_3d_alone = _generate(llm, lora_3d)
96
97 assert len(outputs_2d_alone) == len(VL_TEST_IMAGES)
98 assert len(outputs_3d_alone) == len(VL_TEST_IMAGES)
99 for text in outputs_2d_alone + outputs_3d_alone:
100 assert text, "Empty output from single-adapter LoRA generation"
101
102 # Mixed batch: prompt 0 uses the 2D adapter, prompt 1 uses the 3D
103 # adapter. Per-prompt outputs must match the standalone runs.
104 mixed_outputs = _generate(llm, [lora_2d, lora_3d])
105
106 assert mixed_outputs[0] == outputs_2d_alone[0], (
107 f"Mixed-batch 2D output {mixed_outputs[0]!r} does not match "
108 f"standalone 2D output {outputs_2d_alone[0]!r}"
109 )
110 assert mixed_outputs[1] == outputs_3d_alone[1], (
111 f"Mixed-batch 3D output {mixed_outputs[1]!r} does not match "
112 f"standalone 3D output {outputs_3d_alone[1]!r}"
113 )
114
115 # Reverse assignment: neither adapter should be silently aliased.
116 swapped_outputs = _generate(llm, [lora_3d, lora_2d])
117 assert swapped_outputs[0] == outputs_3d_alone[0], (
118 f"Swapped-batch 3D output {swapped_outputs[0]!r} does not match "
119 f"standalone 3D output {outputs_3d_alone[0]!r}"
120 )
121 assert swapped_outputs[1] == outputs_2d_alone[1], (
122 f"Swapped-batch 2D output {swapped_outputs[1]!r} does not match "
123 f"standalone 2D output {outputs_2d_alone[1]!r}"
124 )
125
126
127 @pytest.mark.skip(reason="This model is too big, so skip this test temporarily.")
128 @pytest.mark.parametrize("fully_sharded_loras", [False, True])
129 @multi_gpu_test(num_gpus=2)
130 def test_qwen36_moe_mixed_2d_3d_lora_tp2(
131 qwen36_moe_2d_lora_files,
132 qwen36_moe_3d_lora_files,
133 fully_sharded_loras,
134 ):
135 _run_mixed_2d_3d_lora_test(
136 lora_2d_files=qwen36_moe_2d_lora_files,
137 lora_3d_files=qwen36_moe_3d_lora_files,
138 tensor_parallel_size=2,
139 fully_sharded_loras=fully_sharded_loras,
140 )
141
142
143 @pytest.mark.skip(reason="This model is too big, so skip this test temporarily.")
144 @pytest.mark.parametrize("fully_sharded_loras", [False, True])
145 @multi_gpu_test(num_gpus=4)
146 def test_qwen36_moe_mixed_2d_3d_lora_tp4(
147 qwen36_moe_2d_lora_files,
148 qwen36_moe_3d_lora_files,
149 fully_sharded_loras,
150 ):
151 _run_mixed_2d_3d_lora_test(
152 lora_2d_files=qwen36_moe_2d_lora_files,
153 lora_3d_files=qwen36_moe_3d_lora_files,
154 tensor_parallel_size=4,
155 fully_sharded_loras=fully_sharded_loras,
156 )