@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 168 lines 4.24 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5
6 from tests.utils import multi_gpu_test
7 from vllm import LLM, PoolingParams
8 from vllm.distributed import cleanup_dist_env_and_memory
9 from vllm.lora.request import LoRARequest
10
11 MODEL_NAME = "Qwen/Qwen3-0.6B"
12 NATIVE_MODEL_NAME = "Skywork/Skywork-Reward-V2-Qwen3-0.6B"
13 PROMPTS = {
14 "star_trek": "Does warp drive appear in Star Trek?",
15 "new_zealand": "Wellington is the capital of New Zealand.",
16 }
17 EXPECTED_OUT = {
18 "star_trek": [13.6640625, -13.6640625],
19 "new_zealand": [-4.80078125, 8.6171875],
20 "native": [-1.52734375],
21 }
22
23
24 def _classify_logits(
25 llm: LLM,
26 prompts: list[str],
27 requests: LoRARequest | list[LoRARequest],
28 ) -> torch.Tensor:
29 params = [PoolingParams(task="classify", use_activation=False) for _ in prompts]
30 outputs = llm.classify(
31 prompts,
32 pooling_params=params,
33 lora_request=requests,
34 use_tqdm=False,
35 )
36 return torch.tensor([output.outputs.probs for output in outputs])
37
38
39 def test_converted_model_with_modules_to_save(
40 qwen3_guard_star_trek_lora_files: str,
41 ) -> None:
42 prompt = PROMPTS["star_trek"]
43 llm = LLM(
44 model=MODEL_NAME,
45 runner="pooling",
46 convert="classify",
47 hf_overrides={"num_labels": 2},
48 dtype="float16",
49 enable_lora=True,
50 max_lora_rank=16,
51 enforce_eager=True,
52 max_model_len=512,
53 gpu_memory_utilization=0.5,
54 )
55 actual = _classify_logits(
56 llm,
57 [prompt],
58 LoRARequest("star-trek", 1, qwen3_guard_star_trek_lora_files),
59 )[0]
60
61 torch.testing.assert_close(
62 actual,
63 torch.tensor(EXPECTED_OUT["star_trek"]),
64 atol=2e-2,
65 rtol=2e-2,
66 )
67
68 del llm
69 cleanup_dist_env_and_memory()
70
71
72 def test_native_classification_model_with_modules_to_save(
73 skywork_qwen3_reward_lora_files: str,
74 ) -> None:
75 prompt = "Which response is more helpful and correct?"
76 llm = LLM(
77 model=NATIVE_MODEL_NAME,
78 runner="pooling",
79 dtype="float16",
80 enable_lora=True,
81 max_lora_rank=8,
82 enforce_eager=True,
83 max_model_len=512,
84 gpu_memory_utilization=0.5,
85 )
86 actual = _classify_logits(
87 llm,
88 [prompt],
89 LoRARequest("native", 1, skywork_qwen3_reward_lora_files),
90 )[0]
91
92 torch.testing.assert_close(
93 actual,
94 torch.tensor(EXPECTED_OUT["native"]),
95 atol=2e-2,
96 rtol=2e-2,
97 )
98
99 del llm
100 cleanup_dist_env_and_memory()
101
102
103 def _test_batched_loras(
104 qwen3_guard_star_trek_lora_files: str,
105 qwen3_guard_new_zealand_lora_files: str,
106 tp_size: int,
107 ) -> None:
108 names = ["star_trek", "new_zealand"]
109 prompts = [PROMPTS[name] for name in names]
110 adapter_paths = [
111 qwen3_guard_star_trek_lora_files,
112 qwen3_guard_new_zealand_lora_files,
113 ]
114 requests = [
115 LoRARequest(name, index, adapter_path)
116 for index, (name, adapter_path) in enumerate(zip(names, adapter_paths), start=1)
117 ]
118
119 llm = LLM(
120 model=MODEL_NAME,
121 runner="pooling",
122 convert="classify",
123 hf_overrides={"num_labels": 1},
124 dtype="float16",
125 enable_lora=True,
126 max_loras=2,
127 max_cpu_loras=2,
128 max_lora_rank=16,
129 max_lora_cls_labels=2,
130 enforce_eager=True,
131 max_model_len=512,
132 gpu_memory_utilization=0.5,
133 tensor_parallel_size=tp_size,
134 )
135 actual = _classify_logits(llm, prompts, requests)
136
137 torch.testing.assert_close(
138 actual,
139 torch.tensor([EXPECTED_OUT[name] for name in names]),
140 atol=2e-2,
141 rtol=2e-2,
142 )
143
144 del llm
145 cleanup_dist_env_and_memory()
146
147
148 def test_batched_loras(
149 qwen3_guard_star_trek_lora_files: str,
150 qwen3_guard_new_zealand_lora_files: str,
151 ) -> None:
152 _test_batched_loras(
153 qwen3_guard_star_trek_lora_files,
154 qwen3_guard_new_zealand_lora_files,
155 tp_size=1,
156 )
157
158
159 @multi_gpu_test(num_gpus=2)
160 def test_batched_loras_tp(
161 qwen3_guard_star_trek_lora_files: str,
162 qwen3_guard_new_zealand_lora_files: str,
163 ) -> None:
164 _test_batched_loras(
165 qwen3_guard_star_trek_lora_files,
166 qwen3_guard_new_zealand_lora_files,
167 tp_size=2,
168 )