@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 195 lines 6.35 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests quantization configuration matching utilities."""
4
5 from unittest.mock import Mock
6
7 import pytest
8
9 from vllm.config.quantization import QuantizationConfigArgs
10 from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
11 should_ignore_layer,
12 )
13 from vllm.model_executor.layers.quantization.online.base import (
14 OnlineQuantizationConfig,
15 _find_matching_targets,
16 )
17 from vllm.model_executor.layers.quantization.quark.utils import (
18 should_ignore_layer as quark_should_ignore_layer,
19 )
20 from vllm.model_executor.layers.quantization.utils.config_utils import (
21 find_matching_patterns,
22 get_layer_name_after_index,
23 is_equal_or_regex_match,
24 )
25
26
27 def test_is_equal_or_regex_match():
28 assert is_equal_or_regex_match(
29 "model.layers.0.mlp.down_proj", "model.layers.0.mlp.down_proj"
30 )
31 assert is_equal_or_regex_match("model.layers.0.mlp.down_proj", r"re:.*down_proj")
32 assert not is_equal_or_regex_match("model.layers.0.mlp.down_proj", "other")
33
34
35 def test_fnmatch_is_opt_in():
36 layer_name = "model.layers.0.mlp.experts"
37 pattern = "*mlp.experts*"
38
39 assert not is_equal_or_regex_match(layer_name, pattern)
40 assert is_equal_or_regex_match(layer_name, pattern, use_fnmatch=True)
41 assert not should_ignore_layer(layer_name, [pattern])
42 assert not quark_should_ignore_layer(layer_name, [pattern])
43 assert should_ignore_layer(layer_name, [pattern], use_fnmatch=True)
44
45
46 def test_online_targets_support_fnmatch_patterns():
47 layer_name = "model.layers.0.mlp.experts"
48 targets = {"*mlp.experts*": "mxfp4"}
49
50 assert _find_matching_targets(layer_name, targets) == ["*mlp.experts*"]
51
52
53 def test_online_ignore_supports_fnmatch_patterns():
54 layer_name = "model.layers.0.mlp.experts"
55 ignore = ["*mlp.experts*"]
56 config = OnlineQuantizationConfig(
57 QuantizationConfigArgs(targets={"*mlp.experts*": "mxfp4"}, ignore=ignore)
58 )
59
60 with pytest.raises(ValueError, match="matches both quantization_config.ignore"):
61 config.resolve_quant_method_cls(Mock(), layer_name)
62
63
64 @pytest.mark.parametrize(
65 "layer_name,expected",
66 [
67 ("model.layers.1.self_attn.qkv_proj", "self_attn.qkv_proj"),
68 ("model.layers.2.mlp.down_proj", "mlp.down_proj"),
69 ("lm_head", "lm_head"),
70 ],
71 )
72 def test_get_layer_name_after_index(layer_name, expected):
73 assert get_layer_name_after_index(layer_name) == expected
74
75
76 @pytest.mark.parametrize(
77 "patterns",
78 [
79 [r"re:.*qkv_proj.*"],
80 [r"re:.*\.[qkv]_proj$"],
81 ],
82 ids=["direct_fused_regex", "fused_shard_regexes"],
83 )
84 def test_find_matching_patterns_for_fused_regexes(patterns):
85 layer_name = "model.layers.0.self_attn.qkv_proj"
86 fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
87
88 matches = find_matching_patterns(layer_name, patterns, fused_mapping)
89 assert all(len(shard_matches) == 1 for shard_matches in matches)
90
91
92 @pytest.mark.parametrize(
93 "patterns,expected",
94 [
95 (
96 [r"re:.*\.qkv_proj$"],
97 [{r"re:.*\.qkv_proj$"}],
98 ),
99 (
100 [r"re:.*\.q_proj$", r"re:.*\.k_proj$", r"re:.*\.v_proj$"],
101 [
102 {r"re:.*\.q_proj$"},
103 {r"re:.*\.k_proj$"},
104 {r"re:.*\.v_proj$"},
105 ],
106 ),
107 ],
108 ids=["direct_fused_match", "individual_shard_matches"],
109 )
110 def test_find_matching_patterns_distinguishes_direct_and_shard_matches(
111 patterns, expected
112 ):
113 layer_name = "model.layers.0.self_attn.qkv_proj"
114 fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
115
116 assert find_matching_patterns(layer_name, patterns, fused_mapping) == expected
117 assert should_ignore_layer(layer_name, patterns, fused_mapping)
118
119
120 def test_should_ignore_layer_rejects_partially_matched_fused_layer():
121 layer_name = "model.layers.0.self_attn.qkv_proj"
122 patterns = [r"re:.*\.q_proj$"]
123 fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
124
125 assert find_matching_patterns(layer_name, patterns, fused_mapping) == [
126 {r"re:.*\.q_proj$"},
127 set(),
128 set(),
129 ]
130 with pytest.raises(ValueError, match="different quantization schemes"):
131 should_ignore_layer(layer_name, patterns, fused_mapping)
132
133
134 def test_should_ignore_layer_returns_false_when_no_fused_pattern_matches():
135 layer_name = "model.layers.0.self_attn.qkv_proj"
136 fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
137
138 assert find_matching_patterns(layer_name, [], fused_mapping) == [
139 set(),
140 set(),
141 set(),
142 ]
143 assert not should_ignore_layer(layer_name, [], fused_mapping)
144
145
146 @pytest.mark.parametrize(
147 "patterns",
148 [
149 [r"re:.*qkv_proj.*"],
150 [r"re:.*\.[qkv]_proj$"],
151 ],
152 ids=["direct_fused_regex", "fused_shard_regexes"],
153 )
154 def test_ignore_and_targets_match_fused_regexes_identically(patterns):
155 layer_name = "model.layers.0.self_attn.qkv_proj"
156 fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
157 targets = {pattern: "fp8_per_block" for pattern in patterns}
158
159 assert should_ignore_layer(layer_name, patterns, fused_mapping)
160 matches = _find_matching_targets(layer_name, targets, fused_mapping)
161 assert len(matches) == 1
162 assert targets[matches[0]] == "fp8_per_block"
163
164
165 def test_ignore_allows_individually_matched_fused_shards():
166 layer_name = "model.layers.0.self_attn.qkv_proj"
167 patterns = [
168 "model.layers.0.self_attn.q_proj",
169 "model.layers.0.self_attn.k_proj",
170 "model.layers.0.self_attn.v_proj",
171 ]
172 fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
173
174 assert should_ignore_layer(layer_name, patterns, fused_mapping)
175
176
177 @pytest.mark.parametrize(
178 "patterns",
179 [
180 ["model.layers.0.self_attn.qkv_proj"],
181 [r"re:.*\.qkv_proj$"],
182 [
183 "model.layers.0.self_attn.q_proj",
184 "model.layers.0.self_attn.k_proj",
185 "model.layers.0.self_attn.v_proj",
186 ],
187 ],
188 ids=["direct_fused_name", "direct_fused_regex", "individual_shards"],
189 )
190 def test_quark_and_compressed_tensors_ignore_fused_layers_identically(patterns):
191 layer_name = "model.layers.0.self_attn.qkv_proj"
192 fused_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
193
194 assert should_ignore_layer(layer_name, patterns, fused_mapping)
195 assert quark_should_ignore_layer(layer_name, patterns, fused_mapping)