@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 343 lines 11.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6
7 import vllm.envs as envs
8 from tests.compile.backend import TestBackend
9 from tests.utils import TestFP8Layer, multi_gpu_test
10 from vllm.compilation.passes.fusion.rms_quant_fusion import RMSNormQuantFusionPass
11 from vllm.compilation.passes.fusion.sequence_parallelism import SequenceParallelismPass
12 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
13 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
14 from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
15 from vllm.config import (
16 CompilationConfig,
17 CUDAGraphMode,
18 DeviceConfig,
19 ModelConfig,
20 PassConfig,
21 VllmConfig,
22 get_current_vllm_config,
23 set_current_vllm_config,
24 )
25 from vllm.config.utils import Range
26 from vllm.distributed import tensor_model_parallel_all_reduce
27 from vllm.distributed.parallel_state import (
28 init_distributed_environment,
29 initialize_model_parallel,
30 )
31 from vllm.model_executor.layers.layernorm import RMSNorm
32 from vllm.model_executor.layers.quantization.utils.quant_utils import (
33 kFp8StaticTensorSym,
34 )
35 from vllm.platforms import current_platform
36 from vllm.utils.system_utils import update_environment_variables
37 from vllm.utils.torch_utils import set_random_seed
38
39 DEVICE_TYPE = current_platform.device_type
40
41 pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
42
43 FP8_DTYPE = current_platform.fp8_dtype()
44 prompts = [
45 "Hello, my name is",
46 "The president of the United States is",
47 "The capital of France is",
48 "The future of AI is",
49 ]
50
51
52 class TestAllReduceRMSNormModel(torch.nn.Module):
53 def __init__(self, hidden_size=16, eps=1e-6):
54 super().__init__()
55 self.hidden_size = hidden_size
56 self.eps = eps
57 self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]
58 self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)]
59
60 def forward(self, x):
61 z = torch.relu(x)
62 x = resid = tensor_model_parallel_all_reduce(z)
63 y = self.norm[0](x)
64
65 z2 = torch.mm(y, self.w[0])
66 x2 = tensor_model_parallel_all_reduce(z2)
67
68 y2, resid = self.norm[1](x2, resid)
69
70 z3 = torch.mm(y2, self.w[1])
71 x3 = tensor_model_parallel_all_reduce(z3)
72
73 y3, resid = self.norm[2](x3, resid)
74
75 z4 = torch.mm(y3, self.w[2])
76 x4 = tensor_model_parallel_all_reduce(z4)
77
78 y4, resid = self.norm[3](x4, resid)
79 return y4
80
81 def ops_in_model_before(self):
82 return [torch.ops.vllm.all_reduce.default]
83
84 def ops_in_model_after(self):
85 return [
86 torch.ops.vllm.all_gather.default,
87 torch.ops.vllm.reduce_scatter.default,
88 ]
89
90 def ops_in_model(self):
91 return [
92 torch.ops.vllm_ir.rms_norm,
93 torch.ops.vllm_ir.fused_add_rms_norm,
94 ]
95
96
97 class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
98 quant_key = kFp8StaticTensorSym
99
100 def __init__(self, hidden_size=16, eps=1e-6):
101 super().__init__()
102 self.vllm_config = get_current_vllm_config()
103 self.hidden_size = hidden_size
104 self.eps = eps
105 self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]
106 self.fp8_linear_layers = [
107 TestFP8Layer(
108 weight_shape=(hidden_size, hidden_size),
109 activation_quant_key=self.quant_key,
110 weight_quant_key=self.quant_key,
111 input_dtype=self.vllm_config.model_config.dtype,
112 )
113 for i in range(3)
114 ]
115
116 def forward(self, hidden_states):
117 # avoid having graph input be an arg to a pattern directly
118 z = torch.relu(hidden_states)
119 x = resid = tensor_model_parallel_all_reduce(z)
120 y = self.norm[0](x)
121
122 z2 = self.fp8_linear_layers[0](y)
123
124 x2 = tensor_model_parallel_all_reduce(z2)
125 y2, resid = self.norm[1](x2, resid)
126
127 z3 = self.fp8_linear_layers[1](y2)
128
129 x3 = tensor_model_parallel_all_reduce(z3)
130 y3, resid = self.norm[2](x3, resid) # use resid here
131
132 z4 = self.fp8_linear_layers[2](y3)
133 x4 = tensor_model_parallel_all_reduce(z4)
134 y4, resid = self.norm[3](x4, resid) # use resid here
135 return y4
136
137 def ops_in_model_after(self):
138 return [
139 torch.ops.vllm.all_gather.default,
140 torch.ops.vllm.reduce_scatter.default,
141 ]
142
143 def ops_in_model_before(self):
144 return [
145 torch.ops.vllm.all_reduce.default,
146 ]
147
148 def ops_in_model(self):
149 if self.vllm_config.compilation_config.pass_config.fuse_norm_quant:
150 return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default]
151 else:
152 quant_ops = (
153 [torch.ops._C.static_scaled_fp8_quant.default]
154 if any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers)
155 else [torch.ops.aten.reciprocal]
156 )
157 return [
158 torch.ops.vllm_ir.rms_norm,
159 torch.ops.vllm_ir.fused_add_rms_norm,
160 *quant_ops,
161 ]
162
163
164 @multi_gpu_test(num_gpus=2)
165 @pytest.mark.parametrize(
166 "test_model_cls, custom_ops",
167 [
168 (TestAllReduceRMSNormModel, "+rms_norm"),
169 (TestAllReduceRMSNormModel, "-rms_norm"),
170 (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,+quant_fp8"),
171 (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,-quant_fp8"),
172 (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,+quant_fp8"),
173 (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,-quant_fp8"),
174 ],
175 )
176 @pytest.mark.parametrize("batch_size", [8])
177 @pytest.mark.parametrize("seq_len", [16])
178 @pytest.mark.parametrize("hidden_size", [16])
179 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
180 @pytest.mark.parametrize("fuse_norm_quant", [True, False])
181 @pytest.mark.parametrize("dynamic", [False, True])
182 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
183 def test_sequence_parallelism_pass(
184 test_model_cls: type[torch.nn.Module],
185 custom_ops: str,
186 batch_size: int,
187 seq_len: int,
188 hidden_size: int,
189 dtype: torch.dtype,
190 fuse_norm_quant: bool,
191 dynamic: bool,
192 ):
193 num_processes = 2
194
195 def run_torch_spawn(fn, nprocs):
196 # need to use torch.mp.spawn otherwise will have problems with
197 # torch.distributed and cuda
198 torch.multiprocessing.spawn(
199 fn,
200 args=(
201 num_processes,
202 test_model_cls,
203 custom_ops,
204 batch_size,
205 seq_len,
206 hidden_size,
207 dtype,
208 fuse_norm_quant,
209 dynamic,
210 ),
211 nprocs=nprocs,
212 )
213
214 run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)
215
216
217 def test_sequence_parallelism_pass_requires_full_graph_compilation():
218 vllm_config = VllmConfig()
219 vllm_config.compilation_config.use_inductor_graph_partition = False
220 vllm_config.compilation_config.splitting_ops = [
221 "vllm::unified_attention_with_output"
222 ]
223
224 sequence_parallelism_pass = object.__new__(SequenceParallelismPass)
225 sequence_parallelism_pass.compilation_config = vllm_config.compilation_config
226 sequence_parallelism_pass.min_token_num = 1
227
228 with pytest.raises(
229 AssertionError,
230 match="SequenceParallelismPass requires full-graph compilation",
231 ):
232 sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8))
233
234
235 def sequence_parallelism_pass_on_test_model(
236 local_rank: int,
237 world_size: int,
238 test_model_cls: type[torch.nn.Module],
239 custom_ops: str,
240 batch_size: int,
241 seq_len: int,
242 hidden_size: int,
243 dtype: torch.dtype,
244 fuse_norm_quant: bool,
245 dynamic: bool,
246 ):
247 set_random_seed(0)
248
249 device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
250 torch.accelerator.set_device_index(device)
251 torch.set_default_device(device)
252 torch.set_default_dtype(dtype)
253
254 update_environment_variables(
255 {
256 "RANK": str(local_rank),
257 "LOCAL_RANK": str(local_rank),
258 "WORLD_SIZE": str(world_size),
259 "MASTER_ADDR": "localhost",
260 "MASTER_PORT": "12345",
261 }
262 )
263
264 # initialize distributed
265 init_distributed_environment()
266
267 # configure vllm config for SequenceParallelismPass
268 custom_ops_list = custom_ops.split(",") if custom_ops else []
269 compilation_config = CompilationConfig(
270 splitting_ops=[], # avoid automatic rms_norm enablement
271 cudagraph_mode=CUDAGraphMode.NONE, # avoid piecewise warnings
272 custom_ops=custom_ops_list,
273 pass_config=PassConfig(
274 enable_sp=True,
275 fuse_norm_quant=fuse_norm_quant,
276 eliminate_noops=True,
277 ),
278 ) # NoOp needed for fusion
279 device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
280
281 # this is a fake model name to construct the model config
282 # in the vllm_config, it's not really used.
283 model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8"
284 model_config = ModelConfig(
285 model=model_name, trust_remote_code=True, dtype=dtype, seed=42
286 )
287
288 vllm_config = VllmConfig(
289 model_config=model_config,
290 device_config=device_config,
291 compilation_config=compilation_config,
292 )
293
294 with set_current_vllm_config(vllm_config):
295 initialize_model_parallel(tensor_model_parallel_size=world_size)
296 noop_pass = NoOpEliminationPass(vllm_config)
297 sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
298 cleanup_pass = PostCleanupPass(vllm_config)
299 assert (
300 sequence_parallelism_pass.compilation_config.splitting_ops
301 == vllm_config.compilation_config.splitting_ops
302 )
303 assert (
304 sequence_parallelism_pass.compilation_config.use_inductor_graph_partition
305 == vllm_config.compilation_config.use_inductor_graph_partition
306 )
307 passes_for_backend: list[VllmInductorPass] = [
308 noop_pass,
309 sequence_parallelism_pass,
310 ]
311
312 if fuse_norm_quant:
313 fusion_pass = RMSNormQuantFusionPass(vllm_config)
314 passes_for_backend.append(fusion_pass)
315
316 passes_for_backend.append(cleanup_pass)
317
318 backend = TestBackend(*passes_for_backend)
319
320 model = test_model_cls(hidden_size)
321
322 hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)
323
324 if dynamic:
325 torch._dynamo.mark_dynamic(hidden_states, 0)
326
327 compiled_model = torch.compile(model, backend=backend)
328 compiled_model(hidden_states)
329
330 assert sequence_parallelism_pass.matched_count == 4
331
332 # In pre-nodes, all reduce should be there,
333 # reduce scatter and all gather should not
334 for op in model.ops_in_model_before():
335 assert backend.op_count(op, before=True) == 4
336
337 # In post-nodes, reduce scatter and all gather should be there,
338 # all reduce should not
339 for op in model.ops_in_model_after():
340 assert backend.op_count(op, before=False) == 4
341
342 for op in model.ops_in_model():
343 assert backend.op_count(op, before=False) > 0