@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 239 lines 8.32 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """Tests for AutoAWQConfig behavior after unification.
5
6 These tests verify the bug fixes for:
7 1. CPU platform override conflict (auto_awq should not override on CPU)
8 2. MoE fallback compatibility (full_config["quant_method"] should be "awq")
9 3. Config attribute consistency
10 4. End-to-end quantization method loading (auto_awq loads and runs correctly)
11
12 Note: Tests that require importing the full auto_awq module (which has GPU-dependent
13 imports) should use subprocess or be run in a GPU environment.
14 """
15
16 from __future__ import annotations
17
18 import pytest
19 import torch
20
21 from tests.quantization.utils import (
22 is_quant_method_supported,
23 load_model_without_vllm_runner,
24 )
25 from vllm.config import set_current_vllm_config
26 from vllm.forward_context import set_forward_context
27 from vllm.model_executor.layers.attention import Attention
28 from vllm.platforms import current_platform
29
30
31 def _get_auto_awq_config_source() -> str:
32 """Read the AutoAWQConfig class source code for isolated testing."""
33 import inspect
34
35 import vllm.model_executor.layers.quantization.auto_awq as auto_awq_module
36
37 return inspect.getsource(auto_awq_module.AutoAWQConfig)
38
39
40 class TestAutoAWQConfigFromConfig:
41 """Tests for AutoAWQConfig.from_config behavior.
42
43 These tests require GPU environment to import the full module.
44 They are skipped on non-GPU platforms.
45 """
46
47 def test_full_config_quant_method_is_awq_for_moe_fallback(self):
48 """full_config should have quant_method='awq' for MoE fallback compatibility.
49
50 MoeWNA16Config only accepts 'gptq' or 'awq' as linear_quant_method.
51 If full_config has 'auto_awq', the MoE fallback will fail.
52 """
53 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
54
55 config = {
56 "w_bit": 4,
57 "q_group_size": 128,
58 "zero_point": True,
59 "lm_head": False,
60 }
61 awq_config = AutoAWQConfig.from_config(config)
62
63 # Verify quant_method is 'awq' for MoE fallback
64 assert awq_config.full_config["quant_method"] == "awq", (
65 f"Expected quant_method='awq', got {awq_config.full_config['quant_method']}"
66 )
67
68 def test_full_config_preserves_other_fields(self):
69 """full_config should preserve all original config fields."""
70 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
71
72 config = {
73 "w_bit": 4,
74 "q_group_size": 128,
75 "zero_point": True,
76 "lm_head": False,
77 "custom_field": "custom_value",
78 }
79 awq_config = AutoAWQConfig.from_config(config)
80
81 assert awq_config.full_config["w_bit"] == 4
82 assert awq_config.full_config["q_group_size"] == 128
83 assert awq_config.full_config["zero_point"] is True
84 assert awq_config.full_config["lm_head"] is False
85 assert awq_config.full_config["custom_field"] == "custom_value"
86
87 def test_full_config_is_copy_not_original(self):
88 """full_config should be a copy, not the original dict."""
89 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
90
91 config = {
92 "w_bit": 4,
93 "q_group_size": 128,
94 "zero_point": True,
95 "lm_head": False,
96 }
97 original_quant_method = config.get("quant_method")
98
99 AutoAWQConfig.from_config(config)
100
101 # Original config should not be modified
102 assert config.get("quant_method") == original_quant_method
103
104
105 class TestAutoAWQConfigAttributes:
106 """Tests for AutoAWQConfig attribute consistency.
107
108 These tests require GPU environment to import the full module.
109 They are skipped on non-GPU platforms.
110 """
111
112 def test_config_attributes_match_input(self):
113 """Config attributes should match input values."""
114 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
115
116 awq_config = AutoAWQConfig(
117 weight_bits=4,
118 group_size=128,
119 zero_point=True,
120 lm_head_quantized=False,
121 modules_to_not_convert=["lm_head"],
122 )
123
124 assert awq_config.weight_bits == 4
125 assert awq_config.group_size == 128
126 assert awq_config.zero_point is True
127 assert awq_config.lm_head_quantized is False
128 assert awq_config.modules_to_not_convert == ["lm_head"]
129
130 def test_pack_factor_for_4bit(self):
131 """Pack factor should be 8 for 4-bit quantization."""
132 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
133
134 awq_config = AutoAWQConfig(
135 weight_bits=4,
136 group_size=128,
137 zero_point=True,
138 lm_head_quantized=False,
139 )
140
141 assert awq_config.pack_factor == 8 # 32 // 4
142
143
144 class TestAutoAWQConfigOverrideLogic:
145 """Tests for override logic by parsing source code (no GPU import required)."""
146
147 def _get_auto_awq_source(self) -> str:
148 """Read the auto_awq.py source file."""
149 import inspect
150 import pathlib
151
152 import vllm.model_executor.layers.quantization.auto_awq as auto_awq_module
153
154 source_path = inspect.getfile(auto_awq_module)
155 return pathlib.Path(source_path).read_text()
156
157 def test_cpu_check_in_override_method(self):
158 """override_quantization_method should check current_platform.is_cpu()."""
159 source = self._get_auto_awq_source()
160
161 # Verify the CPU check exists in override method
162 assert "current_platform.is_cpu()" in source, (
163 "override_quantization_method should check is_cpu()"
164 )
165 assert "return None" in source, (
166 "override_quantization_method should return None on CPU"
167 )
168
169 def test_quant_method_normalization_in_from_config(self):
170 """from_config should normalize quant_method to 'awq' for MoE fallback."""
171 source = self._get_auto_awq_source()
172
173 # Verify the normalization exists
174 assert (
175 '"quant_method"] = "awq"' in source or "'quant_method'] = 'awq'" in source
176 ), "from_config should set quant_method='awq' in full_config"
177
178
179 # =============================================================================
180 # End-to-end integration tests (require GPU environment)
181 # =============================================================================
182
183 PROMPT = "On the surface of Mars, we found"
184
185 # Small AWQ model for testing - using Qwen2 1.5B which has official AWQ checkpoint
186 AWQ_MODELS = [
187 "Qwen/Qwen2-1.5B-Instruct-AWQ",
188 ]
189
190
191 @pytest.mark.skipif(
192 not is_quant_method_supported("auto_awq"),
193 reason="auto_awq is not supported on this GPU type.",
194 )
195 @pytest.mark.parametrize("model_id", AWQ_MODELS)
196 def test_auto_awq_quantization_method(
197 model_id: str, monkeypatch, dist_init, workspace_init
198 ):
199 """Test that quantization='auto_awq' loads and runs correctly."""
200 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
201 model, vllm_config = load_model_without_vllm_runner(
202 model_id,
203 dtype=torch.float16,
204 quantization="auto_awq",
205 model_config_kwargs={"max_model_len": 2048},
206 )
207 target_device = torch.device(current_platform.device_type)
208
209 from vllm.model_executor.layers.quantization.auto_awq import (
210 AutoAWQLinearMethod,
211 AutoAWQMarlinLinearMethod,
212 )
213
214 qkv_proj = model.model.layers[0].self_attn.qkv_proj
215 assert isinstance(
216 qkv_proj.quant_method,
217 (AutoAWQLinearMethod, AutoAWQMarlinLinearMethod),
218 ), (
219 "Expected AutoAWQLinearMethod or AutoAWQMarlinLinearMethod, "
220 f"got {type(qkv_proj.quant_method)}"
221 )
222
223 monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
224 input_ids = torch.tensor([1, 2, 3, 4], device=target_device)
225 positions = torch.arange(input_ids.numel(), device=target_device)
226 with (
227 set_current_vllm_config(vllm_config),
228 set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
229 ):
230 hidden_states = model(input_ids, positions, None)
231 logits = model.compute_logits(hidden_states)
232 assert torch.isfinite(logits).all()
233
234
235 def test_auto_awq_config_get_name():
236 """Test that AutoAWQConfig.get_name() returns 'auto_awq'."""
237 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
238
239 assert AutoAWQConfig.get_name() == "auto_awq"