@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 187 lines 5.75 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import tempfile
5 from pathlib import Path
6 from typing import Any
7
8 import pytest
9 import torch
10
11 from tests.quantization.utils import is_quant_method_supported
12 from vllm import LLM, SamplingParams
13 from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
14 from vllm.platforms import current_platform
15 from vllm.utils.torch_utils import is_torch_equal_or_newer
16
17 from ...utils import create_new_process_for_each_test
18
19
20 def models_list(*, all: bool = True, keywords: list[str] | None = None):
21 TEST_MODELS: list[tuple[str, dict[str, Any]]] = [
22 ("facebook/opt-125m", {}),
23 (
24 "RedHatAI/Llama-3.2-1B-Instruct-FP8-dynamic",
25 {"dtype": torch.float16},
26 ),
27 ]
28
29 if all:
30 if is_quant_method_supported("gptq_marlin"):
31 TEST_MODELS.append(
32 (
33 "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
34 {"quantization": "gptq_marlin"},
35 )
36 )
37
38 if not current_platform.is_rocm() and is_quant_method_supported("awq"):
39 TEST_MODELS.append(
40 ("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {"quantization": "AWQ"})
41 )
42
43 if keywords is None:
44 return TEST_MODELS
45
46 # filter by keywords
47 pred = lambda model: any(keyword in model[0] for keyword in keywords)
48 return list(filter(pred, TEST_MODELS))
49
50
51 @pytest.mark.parametrize(
52 "compilation_mode",
53 [CompilationMode.DYNAMO_TRACE_ONCE, CompilationMode.VLLM_COMPILE],
54 )
55 @pytest.mark.parametrize("model, model_kwargs", models_list(all=True))
56 @create_new_process_for_each_test()
57 def test_full_graph(
58 monkeypatch: pytest.MonkeyPatch,
59 model: str,
60 model_kwargs: dict[str, Any],
61 compilation_mode: CompilationMode,
62 ):
63 with monkeypatch.context():
64 print(f"MODEL={model}")
65
66 run_model(compilation_mode, model, **model_kwargs)
67
68
69 # TODO(luka) add other supported compilation config scenarios here
70 @pytest.mark.parametrize(
71 "compilation_config, model, model_kwargs",
72 [
73 # additional compile sizes, only some of the models
74 (
75 CompilationConfig(mode=CompilationMode.VLLM_COMPILE, compile_sizes=[1, 2]),
76 *model_info,
77 )
78 for model_info in models_list(all=False)
79 ]
80 + [
81 # RMSNorm + quant fusion, only 8-bit quant models
82 (
83 CompilationConfig(
84 mode=CompilationMode.VLLM_COMPILE,
85 custom_ops=["+rms_norm"],
86 pass_config=PassConfig(
87 fuse_norm_quant=True, fuse_act_quant=True, eliminate_noops=True
88 ),
89 ),
90 *model_info,
91 )
92 for model_info in models_list(keywords=["FP8-dynamic"])
93 ]
94 + [
95 # Test depyf integration works
96 (
97 CompilationConfig(
98 mode=CompilationMode.VLLM_COMPILE,
99 debug_dump_path=Path(tempfile.gettempdir()),
100 ),
101 "facebook/opt-125m",
102 {},
103 ),
104 ]
105 + [
106 # graph inductor partition
107 (
108 CompilationConfig(
109 mode=CompilationMode.VLLM_COMPILE,
110 # inductor graph partition uses
111 # torch._C.Tag.cudagraph_unsafe to specify splitting ops
112 use_inductor_graph_partition=True,
113 cudagraph_mode=CUDAGraphMode.PIECEWISE,
114 compile_sizes=[1, 2],
115 ),
116 *model_info,
117 )
118 for model_info in models_list(all=False)
119 if is_torch_equal_or_newer("2.9.0.dev")
120 ]
121 + [
122 # Test get_raw_stream patch with compile_sizes
123 # This tests that TorchInductor autotune works correctly with get_raw_stream
124 # patch in torch 2.9 and without patch in torch 2.10+
125 (
126 CompilationConfig(
127 mode=CompilationMode.VLLM_COMPILE,
128 compile_sizes=[1, 2], # Triggers autotune which uses get_raw_stream
129 cudagraph_mode=CUDAGraphMode.NONE,
130 ),
131 "facebook/opt-125m",
132 {},
133 ),
134 ],
135 )
136 # only test some of the models
137 @create_new_process_for_each_test()
138 def test_custom_compile_config(
139 compilation_config: CompilationConfig,
140 model: str,
141 model_kwargs: dict[str, Any],
142 ):
143 if compilation_config.use_inductor_graph_partition and not is_torch_equal_or_newer(
144 "2.9.0.dev"
145 ):
146 pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
147
148 print(f"MODEL={model}")
149 run_model(compilation_config, model, **model_kwargs)
150
151
152 def run_model(
153 compile_config: CompilationMode | CompilationConfig, model: str, **model_kwargs
154 ):
155 compilation_config = (
156 compile_config
157 if isinstance(compile_config, CompilationConfig)
158 else CompilationConfig(mode=compile_config)
159 )
160
161 prompts = [
162 "Hello, my name is",
163 "The president of the United States is",
164 "The capital of France is",
165 "The future of AI is",
166 ]
167 sampling_params = SamplingParams(temperature=0)
168 # Allow override from model_kwargs
169 model_kwargs = {"tensor_parallel_size": 1, **model_kwargs}
170 model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs}
171
172 # No cudagraphs by default
173 if compilation_config.cudagraph_mode is None:
174 compilation_config.cudagraph_mode = CUDAGraphMode.NONE
175
176 llm = LLM(
177 model=model,
178 compilation_config=compilation_config,
179 **model_kwargs,
180 )
181 outputs = llm.generate(prompts, sampling_params)
182
183 # Print the outputs.
184 for output in outputs:
185 prompt = output.prompt
186 generated_text = output.outputs[0].text
187 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")