@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 370 lines 11.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 from transformers import AutoTokenizer
6
7 import vllm
8 import vllm.config
9 from vllm.assets.image import ImageAsset
10 from vllm.lora.request import LoRARequest
11 from vllm.platforms import current_platform
12
13 from ..utils import create_new_process_for_each_test, multi_gpu_test
14
15 MODEL_PATH = "Qwen/Qwen3.5-4B"
16 TEXT_LORA_ID = 1
17 VL_LORA_ID = 2
18
19 # text-only task
20 TEXT_PROMPT_TEMPLATE = """Write a SQL query for the given database.\nSchema:\nTables:\n - stadium(Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average)\n - singer(Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male)\n - concert(concert_ID, concert_Name, Theme, Stadium_ID, Year)\n - singer_in_concert(concert_ID, Singer_ID)\n\nQuestion:\n{query}""" # noqa: E501
21
22 TEXT_EXPECTED_LORA_OUTPUT = [
23 "SELECT count(*) FROM singer",
24 "SELECT avg(age) , min(age) , max(age) FROM singer WHERE country = 'France'",
25 "SELECT name FROM stadium WHERE stadium_id NOT IN (SELECT stadium_id FROM concert)",
26 ]
27
28
29 # visual caption
30 VL_QUESTION = "What is in the image?"
31 VL_TEST_IMAGES = [
32 ImageAsset("stop_sign"),
33 ImageAsset("cherry_blossom"),
34 ]
35 VL_EXPECTED_LORA_OUTPUT = [
36 'A red STOP sign stands prominently in the foreground, with a traditional Chinese gate adorned with red lanterns and the Chinese characters "中華門" in the background, signaling the entrance to a Chinatown. A black car passes by on the street, and stone lion statues guard the entrance to the culturally rich area.', # noqa: E501
37 "A vibrant blue sky serves as a backdrop for the iconic Tokyo Skytree, partially obscured by the delicate pink blossoms of cherry trees in full bloom.", # noqa: E501
38 ]
39
40 TOKENIZER = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
41
42
43 def _assert_exact_outputs(
44 generated_texts: list[str], expected_outputs: list[str]
45 ) -> None:
46 assert generated_texts == expected_outputs
47
48
49 def _assert_prefix_outputs(
50 generated_texts: list[str],
51 expected_outputs: list[str],
52 ) -> None:
53 assert len(generated_texts) == len(expected_outputs)
54 for generated_text, expected_text in zip(generated_texts, expected_outputs):
55 assert expected_text.startswith(generated_text), (
56 f"Generated {generated_text!r} is not a prefix of expected "
57 f"{expected_text!r}"
58 )
59
60
61 def _run_text_lora_sample(
62 llm: vllm.LLM,
63 lora_path: str,
64 lora_id: int,
65 ) -> list[str]:
66 prompts = [
67 TEXT_PROMPT_TEMPLATE.format(query="How many singers do we have?"),
68 TEXT_PROMPT_TEMPLATE.format(
69 query=(
70 "What is the average, minimum, and maximum "
71 "age of all singers from France?"
72 )
73 ),
74 TEXT_PROMPT_TEMPLATE.format(
75 query="What are the names of the stadiums without any concerts?"
76 ),
77 ]
78 input_templates = []
79 for prompt_text in prompts:
80 messages = [{"role": "user", "content": prompt_text}]
81 prompt = TOKENIZER.apply_chat_template(
82 messages,
83 tokenize=False,
84 add_generation_prompt=True,
85 enable_thinking=False, # disable thinking
86 )
87 input_templates.append(prompt)
88
89 outputs = llm.generate(
90 input_templates,
91 vllm.SamplingParams(temperature=0, max_tokens=512),
92 lora_request=LoRARequest(str(lora_id), lora_id, lora_path),
93 )
94
95 generated_texts: list[str] = []
96 for output in outputs:
97 generated_text = output.outputs[0].text.strip()
98 generated_texts.append(generated_text)
99 print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
100 return generated_texts
101
102
103 def _run_vl_lora_sample(
104 llm: vllm.LLM,
105 lora_path: str | None = None,
106 lora_id: int = 0,
107 ) -> list[str]:
108 messages = [
109 {
110 "role": "user",
111 "content": [
112 {"type": "image"},
113 {"type": "text", "text": VL_QUESTION},
114 ],
115 }
116 ]
117 prompt = TOKENIZER.apply_chat_template(
118 messages,
119 tokenize=False,
120 add_generation_prompt=True,
121 enable_thinking=False,
122 )
123 prompts = [
124 {
125 "prompt": prompt,
126 "multi_modal_data": {"image": asset.pil_image},
127 }
128 for asset in VL_TEST_IMAGES
129 ]
130 outputs = llm.generate(
131 prompts,
132 vllm.SamplingParams(temperature=0, max_tokens=128),
133 lora_request=(
134 LoRARequest(str(lora_id), lora_id, lora_path)
135 if lora_path is not None
136 else None
137 ),
138 )
139
140 generated_texts: list[str] = []
141 for output in outputs:
142 generated_text = output.outputs[0].text.strip()
143 generated_texts.append(generated_text)
144 print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
145 return generated_texts
146
147
148 def _build_text_prompts() -> list[str]:
149 prompts = [
150 TEXT_PROMPT_TEMPLATE.format(query="How many singers do we have?"),
151 TEXT_PROMPT_TEMPLATE.format(
152 query=(
153 "What is the average, minimum, and maximum "
154 "age of all singers from France?"
155 )
156 ),
157 TEXT_PROMPT_TEMPLATE.format(
158 query="What are the names of the stadiums without any concerts?"
159 ),
160 ]
161 input_templates = []
162 for prompt_text in prompts:
163 messages = [{"role": "user", "content": prompt_text}]
164 prompt = TOKENIZER.apply_chat_template(
165 messages,
166 tokenize=False,
167 add_generation_prompt=True,
168 enable_thinking=False,
169 )
170 input_templates.append(prompt)
171 return input_templates
172
173
174 def _build_vl_prompts() -> list[dict]:
175 messages = [
176 {
177 "role": "user",
178 "content": [
179 {"type": "image"},
180 {"type": "text", "text": VL_QUESTION},
181 ],
182 }
183 ]
184 prompt = TOKENIZER.apply_chat_template(
185 messages,
186 tokenize=False,
187 add_generation_prompt=True,
188 enable_thinking=False,
189 )
190 return [
191 {
192 "prompt": prompt,
193 "multi_modal_data": {"image": asset.pil_image},
194 }
195 for asset in VL_TEST_IMAGES
196 ]
197
198
199 def _run_mixed_lora_sample(
200 llm: vllm.LLM,
201 text_lora_path: str,
202 vl_lora_path: str,
203 text_lora_id: int,
204 vl_lora_id: int,
205 ) -> list[str]:
206 text_prompts = _build_text_prompts()[:2]
207 vl_prompts = _build_vl_prompts()
208 prompts = [
209 text_prompts[0],
210 vl_prompts[0],
211 text_prompts[1],
212 vl_prompts[1],
213 ]
214 lora_requests = [
215 LoRARequest("qwen35-text", text_lora_id, text_lora_path),
216 LoRARequest("qwen35-vl", vl_lora_id, vl_lora_path),
217 LoRARequest("qwen35-text", text_lora_id, text_lora_path),
218 LoRARequest("qwen35-vl", vl_lora_id, vl_lora_path),
219 ]
220 outputs = llm.generate(
221 prompts,
222 vllm.SamplingParams(temperature=0, max_tokens=256),
223 lora_request=lora_requests,
224 )
225
226 generated_texts: list[str] = []
227 for output in outputs:
228 generated_text = output.outputs[0].text.strip()
229 generated_texts.append(generated_text)
230 print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
231 return generated_texts
232
233
234 def _run_mixed_lora_and_base_sample(
235 llm: vllm.LLM,
236 text_lora_path: str,
237 vl_lora_path: str,
238 text_lora_id: int,
239 vl_lora_id: int,
240 ) -> list[str]:
241 text_prompt = _build_text_prompts()[0]
242 vl_prompt = _build_vl_prompts()[0]
243 prompts = [
244 text_prompt,
245 vl_prompt,
246 text_prompt,
247 vl_prompt,
248 ]
249 lora_requests = [
250 LoRARequest("qwen35-text", text_lora_id, text_lora_path),
251 LoRARequest("qwen35-vl", vl_lora_id, vl_lora_path),
252 None,
253 None,
254 ]
255 outputs = llm.generate(
256 prompts,
257 vllm.SamplingParams(temperature=0, max_tokens=256),
258 lora_request=lora_requests,
259 )
260
261 generated_texts: list[str] = []
262 for output in outputs:
263 generated_text = output.outputs[0].text.strip()
264 generated_texts.append(generated_text)
265 print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
266 return generated_texts
267
268
269 def _assert_qwen35_text_vl_and_mixed_lora(
270 llm: vllm.LLM,
271 qwen35_text_lora_files: str,
272 qwen35_vl_lora_files: str,
273 ) -> None:
274 generated_texts = _run_text_lora_sample(
275 llm,
276 qwen35_text_lora_files,
277 TEXT_LORA_ID,
278 )
279
280 _assert_exact_outputs(generated_texts, TEXT_EXPECTED_LORA_OUTPUT)
281
282 generated_texts = _run_vl_lora_sample(
283 llm,
284 qwen35_vl_lora_files,
285 VL_LORA_ID,
286 )
287 _assert_prefix_outputs(generated_texts, VL_EXPECTED_LORA_OUTPUT)
288
289 generated_texts = _run_mixed_lora_sample(
290 llm,
291 qwen35_text_lora_files,
292 qwen35_vl_lora_files,
293 text_lora_id=TEXT_LORA_ID,
294 vl_lora_id=VL_LORA_ID,
295 )
296 assert generated_texts[0] == TEXT_EXPECTED_LORA_OUTPUT[0]
297 assert generated_texts[2] == TEXT_EXPECTED_LORA_OUTPUT[1]
298 _assert_prefix_outputs([generated_texts[1]], [VL_EXPECTED_LORA_OUTPUT[0]])
299 _assert_prefix_outputs([generated_texts[3]], [VL_EXPECTED_LORA_OUTPUT[1]])
300
301 generated_texts = _run_mixed_lora_and_base_sample(
302 llm,
303 qwen35_text_lora_files,
304 qwen35_vl_lora_files,
305 text_lora_id=TEXT_LORA_ID,
306 vl_lora_id=VL_LORA_ID,
307 )
308 assert generated_texts[0] == TEXT_EXPECTED_LORA_OUTPUT[0]
309 _assert_prefix_outputs([generated_texts[1]], [VL_EXPECTED_LORA_OUTPUT[0]])
310 assert generated_texts[2] != TEXT_EXPECTED_LORA_OUTPUT[0]
311 assert not VL_EXPECTED_LORA_OUTPUT[0].startswith(generated_texts[3]), (
312 "Non-LoRA vision output unexpectedly matches the LoRA expectation."
313 )
314
315
316 @pytest.mark.skipif(
317 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
318 )
319 @create_new_process_for_each_test()
320 def test_qwen35_text_lora(
321 qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream
322 ):
323 llm = vllm.LLM(
324 model=MODEL_PATH,
325 max_model_len=4096,
326 enable_lora=True,
327 max_loras=2,
328 max_num_seqs=4,
329 max_lora_rank=8,
330 enforce_eager=True,
331 trust_remote_code=True,
332 enable_tower_connector_lora=True,
333 mm_processor_cache_gb=0,
334 limit_mm_per_prompt={"image": 1},
335 )
336
337 _assert_qwen35_text_vl_and_mixed_lora(
338 llm,
339 qwen35_text_lora_files,
340 qwen35_vl_lora_files,
341 )
342
343
344 @multi_gpu_test(num_gpus=4)
345 def test_qwen35_text_lora_tp4(
346 qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream
347 ):
348 llm = vllm.LLM(
349 model=MODEL_PATH,
350 max_model_len=4096,
351 enable_lora=True,
352 max_loras=2,
353 max_lora_rank=8,
354 max_num_seqs=4,
355 enforce_eager=True,
356 tensor_parallel_size=4,
357 trust_remote_code=True,
358 enable_tower_connector_lora=True,
359 mm_processor_cache_gb=0,
360 limit_mm_per_prompt={"image": 1},
361 compilation_config=vllm.config.CompilationConfig(
362 cudagraph_specialize_lora=False,
363 ),
364 )
365
366 _assert_qwen35_text_vl_and_mixed_lora(
367 llm,
368 qwen35_text_lora_files,
369 qwen35_vl_lora_files,
370 )