@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 194 lines 6.15 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests that the auto_gptq quantization method works correctly.
4
5 Run `pytest tests/quantization/test_auto_gptq.py -v -s`.
6 """
7
8 from types import SimpleNamespace
9
10 import pytest
11 import torch
12
13 from tests.quantization.utils import (
14 is_quant_method_supported,
15 load_model_without_vllm_runner,
16 )
17 from vllm.config import set_current_vllm_config
18 from vllm.forward_context import set_forward_context
19 from vllm.model_executor.layers.attention import Attention
20 from vllm.model_executor.layers.fused_moe import RoutedExperts
21 from vllm.model_executor.layers.quantization.auto_gptq import (
22 AutoGPTQConfig,
23 AutoGPTQLinearMethod,
24 AutoGPTQMoEMethod,
25 )
26 from vllm.model_executor.layers.quantization.utils.gptq_utils import override_config
27 from vllm.platforms import current_platform
28
29 PROMPT = "On the surface of Mars, we found"
30
31 MODELS = [
32 "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
33 ]
34
35
36 @pytest.mark.skipif(
37 not is_quant_method_supported("auto_gptq"),
38 reason="auto_gptq is not supported on this GPU type.",
39 )
40 @pytest.mark.parametrize("model_id", MODELS)
41 def test_auto_gptq_quantization_method(
42 model_id: str, monkeypatch, dist_init, workspace_init
43 ):
44 """Test that quantization='auto_gptq' loads and runs correctly."""
45 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
46 model, vllm_config = load_model_without_vllm_runner(
47 model_id,
48 dtype=torch.float16,
49 quantization="auto_gptq",
50 model_config_kwargs={"max_model_len": 2048},
51 )
52
53 qkv_proj = model.model.layers[0].self_attn.qkv_proj
54 assert isinstance(qkv_proj.quant_method, AutoGPTQLinearMethod)
55
56 target_device = torch.device(current_platform.device_type)
57 monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
58 input_ids = torch.tensor([1, 2, 3, 4], device=target_device)
59 positions = torch.arange(input_ids.numel(), device=target_device)
60 with (
61 set_current_vllm_config(vllm_config),
62 set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
63 ):
64 hidden_states = model(input_ids, positions, None)
65 logits = model.compute_logits(hidden_states)
66 assert torch.isfinite(logits).all()
67
68
69 def test_auto_gptq_config_get_name():
70 """Test that AutoGPTQConfig.get_name() returns 'auto_gptq'."""
71 assert AutoGPTQConfig.get_name() == "auto_gptq"
72
73
74 @pytest.mark.parametrize(
75 "dynamic",
76 [
77 {},
78 {r"+:model\.layers\.0\..*": {"desc_act": True}},
79 {
80 r"+:model\.layers\.0\..*": {
81 "desc_act": True,
82 "group_size": 32,
83 }
84 },
85 ],
86 )
87 def test_auto_gptq_rejects_group_activation_order(dynamic):
88 desc_act = not dynamic
89 with pytest.raises(ValueError, match="group activation ordering"):
90 AutoGPTQConfig(4, 128, desc_act, True, False, dynamic, {})
91
92
93 def test_auto_gptq_normalizes_channelwise_activation_order():
94 config = AutoGPTQConfig(4, -1, True, True, False, {}, {})
95 assert not config.desc_act
96
97 config = AutoGPTQConfig(
98 4,
99 128,
100 False,
101 True,
102 False,
103 {r"+:model\.layers\.0\..*": {"desc_act": True, "group_size": -1}},
104 {},
105 )
106 override_config(config, "model.layers.0.self_attn.q_proj")
107 assert config.group_size == -1
108 assert not config.desc_act
109
110 with pytest.raises(ValueError, match="group activation ordering"):
111 AutoGPTQConfig(
112 4,
113 -1,
114 True,
115 True,
116 False,
117 {r"+:model\.layers\.0\..*": {"group_size": 128}},
118 {},
119 )
120
121
122 def test_auto_gptq_moe_creates_zero_initialized_expert_biases():
123 method = object.__new__(AutoGPTQMoEMethod)
124 method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
125 method.input_dtype = None
126 method.experts_cls = None
127 method.moe = SimpleNamespace(w13_num_shards=2)
128 layer = torch.nn.Module()
129
130 method.create_weights(
131 layer=layer,
132 num_experts=2,
133 hidden_size=8,
134 intermediate_size_per_partition=4,
135 params_dtype=torch.float16,
136 weight_loader=lambda *args, **kwargs: None,
137 )
138
139 assert layer.w13_bias.shape == (2, 8)
140 assert layer.w2_bias.shape == (2, 8)
141 assert torch.count_nonzero(layer.w13_bias) == 0
142 assert torch.count_nonzero(layer.w2_bias) == 0
143
144
145 def test_routed_experts_loads_per_expert_biases():
146 class Loader:
147 quant_config = None
148 quant_method = object()
149 moe_config = SimpleNamespace(
150 is_act_and_mul=True,
151 tp_rank=0,
152 moe_parallel_config=SimpleNamespace(tp_size=1),
153 )
154 _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim)
155 _narrow_expert_data_for_padding = staticmethod(
156 RoutedExperts._narrow_expert_data_for_padding
157 )
158 _load_w13 = RoutedExperts._load_w13
159 _loaded_expert_biases = set()
160
161 @staticmethod
162 def _map_global_expert_id_to_local_expert_id(expert_id):
163 return expert_id
164
165 loader = Loader()
166 w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False)
167 w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False)
168
169 for shard_id, loaded in (
170 ("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])),
171 ("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])),
172 ):
173 assert RoutedExperts.weight_loader( # type: ignore[call-overload]
174 loader,
175 w13_bias,
176 loaded,
177 weight_name="model.layers.0.mlp.experts.w13_bias",
178 shard_id=shard_id,
179 expert_id=0,
180 return_success=True,
181 )
182
183 assert RoutedExperts.weight_loader( # type: ignore[call-overload]
184 loader,
185 w2_bias,
186 torch.tensor([9.0, 10.0, 11.0, 12.0]),
187 weight_name="model.layers.0.mlp.experts.w2_bias",
188 shard_id="w2",
189 expert_id=0,
190 return_success=True,
191 )
192 assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8))
193 assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4))
194 assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"}