@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 231 lines 6.14 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 from typing import Literal, NamedTuple
6
7 import pytest
8
9 from vllm.config.model import RunnerOption
10 from vllm.logger import init_logger
11
12 from ..utils import compare_two_settings, create_new_process_for_each_test
13
14 logger = init_logger("test_expert_parallel")
15
16
17 class ParallelSetup(NamedTuple):
18 tp_size: int
19 eager_mode: bool
20 chunked_prefill: bool
21
22
23 class EPTestOptions(NamedTuple):
24 trust_remote_code: bool
25 tokenizer_mode: str | None
26 load_format: str | None = None
27 hf_overrides: str | None = None
28
29
30 @dataclass
31 class EPTestSettings:
32 parallel_setups: list[ParallelSetup]
33 distributed_backends: list[str]
34 runner: RunnerOption
35 test_options: EPTestOptions
36
37 @staticmethod
38 def detailed(
39 *,
40 tp_base: int = 2,
41 runner: RunnerOption = "auto",
42 trust_remote_code: bool = False,
43 tokenizer_mode: str | None = None,
44 load_format: str | None = None,
45 hf_overrides: str | None = None,
46 ):
47 return EPTestSettings(
48 parallel_setups=[
49 ParallelSetup(tp_size=tp_base, eager_mode=False, chunked_prefill=False),
50 ParallelSetup(tp_size=tp_base, eager_mode=False, chunked_prefill=True),
51 ParallelSetup(tp_size=tp_base, eager_mode=True, chunked_prefill=False),
52 ParallelSetup(
53 tp_size=2 * tp_base, eager_mode=False, chunked_prefill=True
54 ),
55 ParallelSetup(
56 tp_size=2 * tp_base, eager_mode=True, chunked_prefill=False
57 ),
58 ],
59 distributed_backends=["mp", "ray"],
60 runner=runner,
61 test_options=EPTestOptions(
62 trust_remote_code=trust_remote_code,
63 tokenizer_mode=tokenizer_mode,
64 load_format=load_format,
65 hf_overrides=hf_overrides,
66 ),
67 )
68
69 @staticmethod
70 def fast(
71 *,
72 tp_base: int = 2,
73 runner: RunnerOption = "auto",
74 trust_remote_code: bool = False,
75 tokenizer_mode: str | None = None,
76 load_format: str | None = None,
77 hf_overrides: str | None = None,
78 ):
79 return EPTestSettings(
80 parallel_setups=[
81 ParallelSetup(tp_size=tp_base, eager_mode=True, chunked_prefill=False),
82 ],
83 distributed_backends=["mp"],
84 runner=runner,
85 test_options=EPTestOptions(
86 trust_remote_code=trust_remote_code,
87 tokenizer_mode=tokenizer_mode,
88 load_format=load_format,
89 hf_overrides=hf_overrides,
90 ),
91 )
92
93 def iter_params(self, model_name: str):
94 opts = self.test_options
95
96 for parallel_setup in self.parallel_setups:
97 for distributed_backend in self.distributed_backends:
98 yield (
99 model_name,
100 parallel_setup,
101 distributed_backend,
102 self.runner,
103 opts,
104 )
105
106
107 # NOTE: You can adjust tp_base locally to fit the model in GPU
108 # The values displayed here are only a rough indicator of the size of the model
109
110 TEST_MODELS = {
111 "deepseek-ai/DeepSeek-V2-Lite-Chat": EPTestSettings.fast(trust_remote_code=True),
112 "mistralai/Mixtral-8x7B-Instruct-v0.1": EPTestSettings.fast(tp_base=4),
113 }
114
115
116 def _compare_tp(
117 model_name: str,
118 parallel_setup: ParallelSetup,
119 distributed_backend: str,
120 runner: RunnerOption,
121 test_options: EPTestOptions,
122 num_gpus_available: int,
123 *,
124 method: Literal["generate"],
125 ):
126 (
127 tp_size,
128 eager_mode,
129 chunked_prefill,
130 ) = parallel_setup
131 (
132 trust_remote_code,
133 tokenizer_mode,
134 load_format,
135 hf_overrides,
136 ) = test_options
137
138 if num_gpus_available < tp_size:
139 pytest.skip(f"Need at least {tp_size} GPUs")
140
141 common_args = [
142 # use half precision for speed and memory savings in CI environment
143 "--dtype",
144 "float16",
145 "--max-model-len",
146 "2048",
147 "--max-num-seqs",
148 "8",
149 "--load-format",
150 "auto",
151 ]
152 if chunked_prefill:
153 common_args.append("--enable-chunked-prefill")
154 if eager_mode:
155 common_args.append("--enforce-eager")
156 if runner != "auto":
157 common_args.extend(["--runner", runner])
158 if trust_remote_code:
159 common_args.append("--trust-remote-code")
160 if tokenizer_mode:
161 common_args.extend(["--tokenizer-mode", tokenizer_mode])
162 if load_format:
163 common_args.extend(["--load-format", load_format])
164 if hf_overrides:
165 common_args.extend(["--hf-overrides", hf_overrides])
166
167 ep_env = {
168 "VLLM_TEST_ENABLE_EP": "1",
169 }
170
171 ep_args = [
172 *common_args,
173 "--tensor-parallel-size",
174 str(tp_size),
175 "--distributed-executor-backend",
176 distributed_backend,
177 ]
178
179 # compare without expert parallelism
180 tp_env = {
181 "VLLM_TEST_ENABLE_EP": "0",
182 }
183
184 tp_args = [
185 *common_args,
186 "--tensor-parallel-size",
187 str(tp_size),
188 "--distributed-executor-backend",
189 "mp",
190 ]
191
192 try:
193 compare_two_settings(
194 model_name,
195 ep_args,
196 tp_args,
197 ep_env,
198 tp_env,
199 method=method,
200 max_wait_seconds=360,
201 )
202 except Exception:
203 raise
204
205
206 @pytest.mark.parametrize(
207 ("model_name", "parallel_setup", "distributed_backend", "runner", "test_options"),
208 [
209 params
210 for model_name, settings in TEST_MODELS.items()
211 for params in settings.iter_params(model_name)
212 ],
213 )
214 @create_new_process_for_each_test()
215 def test_ep(
216 model_name: str,
217 parallel_setup: ParallelSetup,
218 distributed_backend: str,
219 runner: RunnerOption,
220 test_options: EPTestOptions,
221 num_gpus_available,
222 ):
223 _compare_tp(
224 model_name,
225 parallel_setup,
226 distributed_backend,
227 runner,
228 test_options,
229 num_gpus_available,
230 method="generate",
231 )