@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 227 lines 7.96 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 import shutil
6 from collections.abc import Sequence
7
8 import pytest
9 import torch
10 from safetensors.torch import load_file, save_file
11
12 import vllm
13 from vllm.lora.request import LoRARequest
14 from vllm.platforms import current_platform
15
16 from ..utils import multi_gpu_test
17
18 MODEL_PATH = "allenai/OLMoE-1B-7B-0125-Instruct"
19
20 PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me. Do not return any additional explanation. Below is an instruction that describes a task, Write a response that appropriately completes the request.
21 "
22 ##Instruction:
23 candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
24 Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
25 The People_ID of candidate is the foreign key of People_ID of people.
26
27
28 ###Input:
29 {context}
30
31 ###Response:""" # noqa: E501
32
33 EXPECTED_LORA_OUTPUT = [
34 "SELECT count(*) FROM candidate",
35 "SELECT count(*) FROM candidate",
36 "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
37 "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
38 ]
39
40 EXPECTED_BASE_MODEL_OUTPUT = [
41 "SELECT COUNT(Candidate_ID) FROM candidate",
42 "SELECT COUNT(Candidate_ID) FROM candidate",
43 "SELECT Candidate_ID, COUNT(*) as Total_Candidates\nFROM candidate\nINNER JOIN people ON candidate.People_ID = people.People_ID", # noqa: E501
44 # There are multiple acceptable responses
45 (
46 "SELECT Candidate_ID, Poll_Source FROM candidate WHERE People_ID IN (SELECT People_ID FROM people) ORDER BY COUNT(*) DESC LIMIT 1", # noqa: E501
47 "SELECT Candidate_ID, Poll_Source FROM candidate WHERE COUNT(People_ID) = (SELECT COUNT(People_ID) FROM people) ORDER BY Candidate_ID DESC LIMIT 1", # noqa: E501
48 ),
49 ]
50
51
52 def _output_matches(generated: str, accepted: str | Sequence[str]) -> bool:
53 if isinstance(accepted, str):
54 accepted = (accepted,)
55 return any(generated.startswith(s) for s in accepted)
56
57
58 def generate_and_test(
59 llm: vllm.LLM,
60 lora_path: str,
61 lora_id: list[int | None] | int | None,
62 compare_lower: bool = False,
63 ) -> None:
64 prompts = [
65 PROMPT_TEMPLATE.format(context="How many candidates are there?"),
66 PROMPT_TEMPLATE.format(context="Count the number of candidates."),
67 PROMPT_TEMPLATE.format(
68 context="Which poll resource provided the most number of candidate information?" # noqa: E501
69 ),
70 PROMPT_TEMPLATE.format(
71 context="Return the poll resource associated with the most candidates."
72 ),
73 ]
74
75 lora_request = None
76 if isinstance(lora_id, int):
77 lora_request = LoRARequest(str(lora_id), lora_id, lora_path)
78 elif isinstance(lora_id, list):
79 lora_request = [
80 LoRARequest(str(i), i, lora_path) if i is not None else None
81 for i in lora_id
82 ]
83
84 sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
85 outputs = llm.generate(prompts, sampling_params, lora_request=lora_request)
86 # Print the outputs.
87 generated_texts: list[str] = []
88 for output in outputs:
89 prompt = output.prompt
90 generated_text = output.outputs[0].text.strip()
91 generated_texts.append(generated_text)
92 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
93
94 for i in range(len(EXPECTED_LORA_OUTPUT)):
95 req_lora_id = lora_id[i] if isinstance(lora_id, list) else lora_id
96 generated_text = generated_texts[i]
97 expected_output = (
98 EXPECTED_LORA_OUTPUT[i]
99 if req_lora_id is not None
100 else EXPECTED_BASE_MODEL_OUTPUT[i]
101 )
102
103 if compare_lower:
104 generated_text = generated_text.lower()
105 if isinstance(expected_output, str):
106 expected_output = (expected_output.lower(),)
107 else:
108 expected_output = tuple(s.lower() for s in expected_output)
109 assert _output_matches(generated_text, expected_output), (
110 f"Output {i}: {generated_text!r} does not match any of {expected_output!r}"
111 )
112
113
114 @pytest.mark.skipif(
115 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
116 )
117 def test_olmoe_lora(olmoe_lora_files, maybe_enable_lora_dual_stream):
118 # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
119 # Otherwise, the lora-test will fail due to CUDA OOM.
120 llm = vllm.LLM(
121 MODEL_PATH,
122 max_model_len=1024,
123 enable_lora=True,
124 max_loras=4,
125 enforce_eager=True,
126 trust_remote_code=True,
127 enable_chunked_prefill=True,
128 )
129
130 generate_and_test(llm, olmoe_lora_files, lora_id=1)
131 generate_and_test(llm, olmoe_lora_files, lora_id=2)
132
133
134 def test_olmoe_lora_mixed(olmoe_lora_files):
135 llm = vllm.LLM(
136 MODEL_PATH,
137 max_model_len=1024,
138 enable_lora=True,
139 max_loras=4,
140 enforce_eager=True,
141 trust_remote_code=True,
142 hf_overrides={"head_dtype": "float32"},
143 enable_chunked_prefill=True,
144 )
145
146 generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])
147
148
149 def test_olmoe_lora_mixed_random(
150 olmoe_lora_files, tmp_path, maybe_enable_lora_dual_stream
151 ):
152 # Create a dummy LoRA with random weights based on the real one
153 random_lora_path = tmp_path / "random_lora"
154 shutil.copytree(olmoe_lora_files, random_lora_path)
155
156 weights_path = random_lora_path / "adapter_model.safetensors"
157 weights = load_file(str(weights_path))
158 random_weights = {k: torch.randn_like(v) for k, v in weights.items()}
159 save_file(random_weights, str(weights_path))
160
161 llm = vllm.LLM(
162 MODEL_PATH,
163 max_model_len=1024,
164 enable_lora=True,
165 max_loras=4,
166 enforce_eager=True,
167 trust_remote_code=True,
168 enable_chunked_prefill=True,
169 )
170
171 prompts = [
172 PROMPT_TEMPLATE.format(context="How many candidates are there?"),
173 PROMPT_TEMPLATE.format(context="Count the number of candidates."),
174 ]
175
176 lora_requests = [
177 LoRARequest("real", 1, olmoe_lora_files),
178 LoRARequest("random", 2, str(random_lora_path)),
179 ]
180
181 sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
182 outputs = llm.generate(prompts, sampling_params, lora_request=lora_requests)
183 assert outputs[0].outputs[0].text.strip().startswith(EXPECTED_LORA_OUTPUT[0])
184
185
186 @pytest.mark.skipif(
187 current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
188 )
189 @pytest.mark.parametrize("fully_sharded_loras", [False, True])
190 @multi_gpu_test(num_gpus=2)
191 def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
192 llm = vllm.LLM(
193 MODEL_PATH,
194 max_model_len=1024,
195 enable_lora=True,
196 max_loras=4,
197 enforce_eager=True,
198 trust_remote_code=True,
199 enable_chunked_prefill=True,
200 tensor_parallel_size=2,
201 fully_sharded_loras=fully_sharded_loras,
202 )
203
204 generate_and_test(llm, olmoe_lora_files, lora_id=1)
205 generate_and_test(llm, olmoe_lora_files, lora_id=2)
206
207
208 @pytest.mark.parametrize("fully_sharded_loras", [False, True])
209 @multi_gpu_test(num_gpus=4)
210 def test_olmoe_lora_tp4(olmoe_lora_files, fully_sharded_loras):
211 llm = vllm.LLM(
212 MODEL_PATH,
213 max_model_len=1024,
214 enable_lora=True,
215 max_loras=4,
216 enforce_eager=True,
217 trust_remote_code=True,
218 enable_chunked_prefill=True,
219 tensor_parallel_size=4,
220 fully_sharded_loras=fully_sharded_loras,
221 )
222 generate_and_test(
223 llm, olmoe_lora_files, lora_id=1, compare_lower=fully_sharded_loras
224 )
225 generate_and_test(
226 llm, olmoe_lora_files, lora_id=2, compare_lower=fully_sharded_loras
227 )