Skip to content

Commit 7edb10a

Browse files
tianleiwutianleiwu
committed
format
1 parent 5dde8e7 commit 7edb10a

2 files changed

Lines changed: 22 additions & 122 deletions

File tree

‎src/python/py/models/builders/cuda_quantizer.py‎

Lines changed: 1 addition & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -97,40 +97,7 @@ def _preprocess_weights_for_mixed_gemm_torch(tensor, bits: int, sm: int):
9797

9898
permutation_map = {
9999
"16_8": [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15],
100-
"16_4": [
101-
0,
102-
1,
103-
8,
104-
9,
105-
16,
106-
17,
107-
24,
108-
25,
109-
2,
110-
3,
111-
10,
112-
11,
113-
18,
114-
19,
115-
26,
116-
27,
117-
4,
118-
5,
119-
12,
120-
13,
121-
20,
122-
21,
123-
28,
124-
29,
125-
6,
126-
7,
127-
14,
128-
15,
129-
22,
130-
23,
131-
30,
132-
31,
133-
],
100+
"16_4": [0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31], # fmt: skip
134101
}
135102
mma_shape_n = 8
136103
b_rows_per_mma = 8 * 16 // bits_b

‎test/python/builder/test_tied_embeddings.py‎

Lines changed: 21 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -222,96 +222,29 @@ def test_shared_embeddings_prefers_quantized_path_only_when_both_layers_are_quan
222222
assert model.tied_unquantized_embeddings is expected_tied_unquantized
223223

224224

225+
# fmt: off
226+
_TIED_QUANTIZED_EMBEDDING_WEIGHT_NAME_CASES = [
227+
("default", 32, True, 4, "lm_head.MatMul.weight_Q4", "lm_head.MatMul.weight_scales", ""),
228+
("default", 32, False, 4, "lm_head.MatMul.weight", "", ""),
229+
("rtn", 32, True, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", ""),
230+
("rtn", 32, False, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
231+
("rtn_last", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", ""),
232+
("rtn_last", 32, False, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
233+
("k_quant", 32, True, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
234+
("k_quant", 32, False, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
235+
("k_quant_last", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
236+
("k_quant_last", 32, False, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
237+
("k_quant_mixed", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
238+
("k_quant_mixed", 32, False, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
239+
("k_quant_linear", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
240+
("k_quant_linear", 32, False, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
241+
]
242+
# fmt: on
243+
244+
225245
@pytest.mark.parametrize(
226246
"int4_algo_config, matmul_block_size, is_symmetric, expected_bits, expected_weight, expected_scale, expected_zp",
227-
[
228-
("default", 32, True, 4, "lm_head.MatMul.weight_Q4", "lm_head.MatMul.weight_scales", ""),
229-
("default", 32, False, 4, "lm_head.MatMul.weight", "", ""),
230-
("rtn", 32, True, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", ""),
231-
("rtn", 32, False, 4, "lm_head.MatMul.weight_Q4G32", "lm_head.MatMul.weight_scale", "lm_head.MatMul.weight_zp"),
232-
("rtn_last", 32, True, 8, "lm_head.MatMul.weight_Q8G32", "lm_head.MatMul.weight_scale", ""),
233-
(
234-
"rtn_last",
235-
32,
236-
False,
237-
8,
238-
"lm_head.MatMul.weight_Q8G32",
239-
"lm_head.MatMul.weight_scale",
240-
"lm_head.MatMul.weight_zp",
241-
),
242-
(
243-
"k_quant",
244-
32,
245-
True,
246-
4,
247-
"lm_head.MatMul.weight_Q4G32",
248-
"lm_head.MatMul.weight_scale",
249-
"lm_head.MatMul.weight_zp",
250-
),
251-
(
252-
"k_quant",
253-
32,
254-
False,
255-
4,
256-
"lm_head.MatMul.weight_Q4G32",
257-
"lm_head.MatMul.weight_scale",
258-
"lm_head.MatMul.weight_zp",
259-
),
260-
(
261-
"k_quant_last",
262-
32,
263-
True,
264-
8,
265-
"lm_head.MatMul.weight_Q8G32",
266-
"lm_head.MatMul.weight_scale",
267-
"lm_head.MatMul.weight_zp",
268-
),
269-
(
270-
"k_quant_last",
271-
32,
272-
False,
273-
8,
274-
"lm_head.MatMul.weight_Q8G32",
275-
"lm_head.MatMul.weight_scale",
276-
"lm_head.MatMul.weight_zp",
277-
),
278-
(
279-
"k_quant_mixed",
280-
32,
281-
True,
282-
8,
283-
"lm_head.MatMul.weight_Q8G32",
284-
"lm_head.MatMul.weight_scale",
285-
"lm_head.MatMul.weight_zp",
286-
),
287-
(
288-
"k_quant_mixed",
289-
32,
290-
False,
291-
8,
292-
"lm_head.MatMul.weight_Q8G32",
293-
"lm_head.MatMul.weight_scale",
294-
"lm_head.MatMul.weight_zp",
295-
),
296-
(
297-
"k_quant_linear",
298-
32,
299-
True,
300-
8,
301-
"lm_head.MatMul.weight_Q8G32",
302-
"lm_head.MatMul.weight_scale",
303-
"lm_head.MatMul.weight_zp",
304-
),
305-
(
306-
"k_quant_linear",
307-
32,
308-
False,
309-
8,
310-
"lm_head.MatMul.weight_Q8G32",
311-
"lm_head.MatMul.weight_scale",
312-
"lm_head.MatMul.weight_zp",
313-
),
314-
],
247+
_TIED_QUANTIZED_EMBEDDING_WEIGHT_NAME_CASES,
315248
)
316249
def test_tied_quantized_embedding_weight_names_cover_all_supported_algorithms(
317250
int4_algo_config,

0 commit comments

Comments
 (0)