@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 303 lines 9.59 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """This script contains:
4 1. test multi loras service with tp >= 2
5 2. test multi loras request
6 """
7
8 import os
9
10 import pytest
11
12 from tests.utils import multi_gpu_test
13 from vllm import LLM, SamplingParams
14 from vllm.exceptions import VLLMValidationError
15 from vllm.lora.request import LoRARequest
16
17 MODEL_PATH = "Qwen/Qwen3-0.6B"
18 LORA_NAME_PATH_MAP = {
19 "Alice": "charent/self_cognition_Alice",
20 "Bob": "charent/self_cognition_Bob",
21 "Cat": "charent/self_cognition_Bob", # same as Bob
22 }
23
24 LORA_NAME_ID_MAP = {}
25 INCREASE_LORA_ID = 0
26 LORA_RANK = 8
27
28 LORA_TEST_PROMPTS = ["What is GitHub?", "Hi, tell me about you"]
29 LORA_TEST_EXPECTED = [
30 "GitHub is an open-source platform that provides a way to manage and develop software projects. It allows developers to store and manage code, collaborate on projects, and automate tasks.", # noqa: E501
31 "I am Alice, an AI assistant developed by GitHub/Charent.",
32 ]
33
34
35 def format_chatml_messages(
36 prompt: str, system_prompt: str = "You are a helpful assistant."
37 ) -> list[dict[str, str]]:
38 return [
39 {"role": "system", "content": system_prompt},
40 {"role": "user", "content": prompt},
41 ]
42
43
44 @pytest.fixture(autouse=True)
45 def set_mrv2_env():
46 original = os.environ.get("VLLM_USE_V2_MODEL_RUNNER", "0")
47 os.environ["VLLM_USE_V2_MODEL_RUNNER"] = "1"
48 yield
49
50 if original is None:
51 os.environ.pop("VLLM_USE_V2_MODEL_RUNNER", None)
52 else:
53 os.environ["VLLM_USE_V2_MODEL_RUNNER"] = original
54
55
56 def make_add_lora_request(name: str, path: str):
57 global INCREASE_LORA_ID, LORA_NAME_ID_MAP
58
59 INCREASE_LORA_ID += 1
60 LORA_NAME_ID_MAP[name] = INCREASE_LORA_ID
61
62 return LoRARequest(
63 lora_name=name,
64 lora_int_id=INCREASE_LORA_ID,
65 lora_path=path,
66 )
67
68
69 @multi_gpu_test(num_gpus=2)
70 def test_multi_loras_with_tp_sync():
71 llm = LLM(
72 model=MODEL_PATH,
73 enable_lora=True,
74 max_loras=2, # ensure max_loras < max_cpu_loras
75 max_lora_rank=LORA_RANK,
76 max_model_len=512,
77 gpu_memory_utilization=0.5,
78 tensor_parallel_size=2, # ensure tp >= 2
79 max_cpu_loras=4, # ensure max_cpu_loras >= 2
80 )
81
82 def run_check_lora(fn, args, expected: list):
83 fn(args)
84 assert set(llm.llm_engine.list_loras()) == set(expected)
85
86 # simulate add loras with CLI args
87 # likes: `--lora-modules Alice=/path/to/Alice Bob=/path/to/Bob`
88 run_check_lora(
89 llm.llm_engine.add_lora,
90 make_add_lora_request("Alice", LORA_NAME_PATH_MAP["Alice"]),
91 [1],
92 )
93 run_check_lora(
94 llm.llm_engine.add_lora,
95 make_add_lora_request("Bob", LORA_NAME_PATH_MAP["Bob"]),
96 [1, 2],
97 )
98 run_check_lora(
99 llm.llm_engine.add_lora,
100 make_add_lora_request("Cat", LORA_NAME_PATH_MAP["Cat"]),
101 [1, 2, 3],
102 )
103
104 # set temperature = 0 for greedy search
105 sampling_params = SamplingParams(temperature=0, max_tokens=64)
106
107 def call_llm_get_outputs(prompt: str, lora_name: str):
108 lora_request = LoRARequest(
109 lora_name=lora_name,
110 lora_int_id=LORA_NAME_ID_MAP[lora_name],
111 lora_path=LORA_NAME_PATH_MAP[lora_name],
112 )
113 messages = format_chatml_messages(prompt)
114 outputs = llm.chat(
115 [messages],
116 sampling_params,
117 chat_template_kwargs={
118 "enable_thinking": False
119 }, # for those loras, ensure enable_thinking=False
120 lora_request=lora_request,
121 use_tqdm=False,
122 )
123 output_text = outputs[0].outputs[0].text
124 return output_text
125
126 def reload_lora(name: str):
127 """Reload a lora to simulate the case:
128 setting `VLLM_ALLOW_RUNTIME_LORA_UPDATING=true`
129 for dynamic lora loading and unloading
130 """
131 remove_lora_response = llm.llm_engine.remove_lora(
132 lora_id=LORA_NAME_ID_MAP[name]
133 )
134
135 add_lora_response = llm.llm_engine.add_lora(
136 make_add_lora_request(name, LORA_NAME_PATH_MAP[name])
137 )
138
139 print(f"{remove_lora_response=}, {add_lora_response=}")
140
141 def check_outputs(outputs: str, expected: str):
142 print(f"{prompt=}.\n{expected_output=}\n{output_text=}")
143 print("\n----------------------------\n")
144 assert outputs == expected
145
146 for prompt, expected_output in zip(LORA_TEST_PROMPTS, LORA_TEST_EXPECTED):
147 output_text = call_llm_get_outputs(prompt, "Alice")
148 check_outputs(output_text, expected_output)
149
150 # call Bob, ignore what it is output
151 call_llm_get_outputs(prompt, "Bob")
152 print("After call Bob:")
153
154 # call Alice
155 output_text = call_llm_get_outputs(prompt, "Alice")
156 check_outputs(output_text, expected_output)
157
158 # reload Bob Lora
159 reload_lora("Bob")
160 print("After reload Bob:")
161
162 # call Alice
163 output_text = call_llm_get_outputs(prompt, "Alice")
164 check_outputs(output_text, expected_output)
165
166 # reload Alice Lora
167 reload_lora("Alice")
168 print("After reload Alice:")
169
170 output_text = call_llm_get_outputs(prompt, "Alice")
171 check_outputs(output_text, expected_output)
172
173
174 def test_multiple_lora_requests():
175 llm = LLM(
176 model=MODEL_PATH,
177 enable_lora=True,
178 max_loras=4,
179 max_lora_rank=LORA_RANK,
180 max_model_len=512,
181 gpu_memory_utilization=0.5,
182 )
183 PROMPTS = ["Hello, my name is"] * 2
184 LORA_NAME = "Alice"
185 lora_request = [
186 LoRARequest(LORA_NAME + str(idx), idx + 1, LORA_NAME_PATH_MAP[LORA_NAME])
187 for idx in range(len(PROMPTS))
188 ]
189 # Multiple SamplingParams should be matched with each prompt
190 outputs = llm.generate(PROMPTS, lora_request=lora_request)
191 assert len(PROMPTS) == len(outputs)
192
193 # Exception raised, if the size of params does not match the size of prompts
194 with pytest.raises(VLLMValidationError):
195 outputs = llm.generate(PROMPTS, lora_request=lora_request[:1])
196
197 # Single LoRARequest should be applied to every prompt
198 single_lora_request = lora_request[0]
199 outputs = llm.generate(PROMPTS, lora_request=single_lora_request)
200 assert len(PROMPTS) == len(outputs)
201
202
203 def test_load_inplace_offline_reload(
204 qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str
205 ) -> None:
206 """Test that load_inplace=True allows reloading LoRA adapters with the same ID
207 in offline mode (using LLM class directly).
208 """
209 llm = LLM(
210 model=MODEL_PATH,
211 enable_lora=True,
212 max_loras=2,
213 max_lora_rank=LORA_RANK,
214 max_model_len=512,
215 gpu_memory_utilization=0.5,
216 )
217 adapter_id = 1
218 messages = format_chatml_messages(
219 "Make your favorite animal noise.",
220 system_prompt="Follow the instructions to make animal noises",
221 )
222 sampling_params = SamplingParams(temperature=0, max_tokens=10)
223
224 # Load meowing LoRA with load_inplace=True
225 meowing_request = LoRARequest(
226 lora_name="test-adapter",
227 lora_int_id=adapter_id,
228 lora_path=qwen3_meowing_lora_files,
229 )
230
231 outputs = llm.chat([messages], sampling_params, lora_request=meowing_request)
232 first_output = outputs[0].outputs[0].text.strip()
233 assert "Meow Meow Meow" in first_output, (
234 f"Expected meowing output, got: {first_output}"
235 )
236
237 # Reload with woofing LoRA (same ID, different weights, load_inplace=True)
238 woofing_request = LoRARequest(
239 lora_name="test-adapter-woof",
240 lora_int_id=adapter_id, # Same ID
241 lora_path=qwen3_woofing_lora_files, # Different weights
242 load_inplace=True, # Force reload
243 )
244
245 outputs = llm.chat([messages], sampling_params, lora_request=woofing_request)
246 second_output = outputs[0].outputs[0].text.strip()
247 assert "Woof Woof Woof" in second_output, (
248 f"Expected woofing output, got: {second_output}"
249 )
250
251
252 def test_load_inplace_false_no_reload(
253 qwen3_meowing_lora_files: str, qwen3_woofing_lora_files: str
254 ) -> None:
255 """Test that load_inplace=False prevents reloading when an adapter
256 with the same ID already exists.
257 """
258 llm = LLM(
259 model=MODEL_PATH,
260 enable_lora=True,
261 max_loras=2,
262 max_lora_rank=LORA_RANK,
263 max_model_len=512,
264 gpu_memory_utilization=0.5,
265 )
266 adapter_id = 2
267 messages = format_chatml_messages(
268 "Make your favorite animal noise.",
269 system_prompt="Follow the instructions to make animal noises",
270 )
271 sampling_params = SamplingParams(temperature=0, max_tokens=10)
272
273 # Load meowing LoRA first with load_inplace=True
274 meowing_request_initial = LoRARequest(
275 lora_name="test-adapter-2",
276 lora_int_id=adapter_id,
277 lora_path=qwen3_meowing_lora_files,
278 )
279
280 outputs = llm.chat(
281 [messages], sampling_params, lora_request=meowing_request_initial
282 )
283 first_output = outputs[0].outputs[0].text.strip()
284 assert "Meow Meow Meow" in first_output, (
285 f"Expected meowing output, got: {first_output}"
286 )
287
288 # Try to load woofing LoRA with same ID but load_inplace=False
289 # This should NOT reload (adapter 2 already exists)
290 woofing_request_no_reload = LoRARequest(
291 lora_name="test-adapter-2-woof",
292 lora_int_id=adapter_id, # Same ID
293 lora_path=qwen3_woofing_lora_files,
294 )
295
296 outputs = llm.chat(
297 [messages], sampling_params, lora_request=woofing_request_no_reload
298 )
299 second_output = outputs[0].outputs[0].text.strip()
300 # Should still get meowing output because it didn't reload
301 assert "Meow Meow Meow" in second_output, (
302 f"Expected meowing output (no reload), got: {second_output}"
303 )