Skip to content

Commit 93c0d54

Browse files
committed
Fix proj linear count
1 parent 415c0b7 commit 93c0d54

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

tests/test_auto_fp8.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
MODELS = [
1111
("facebook/opt-125m", 160),
12-
("Qwen/Qwen2-0.5B-Instruct", 600),
12+
("Qwen/Qwen2-0.5B-Instruct", 620),
1313
]
1414

1515
@pytest.mark.parametrize("model_id,target_size", MODELS)
@@ -83,7 +83,7 @@ def test_kv_cache_static_quantization(model_id, target_size):
8383
proj_linear_count = 0
8484
output_scale_count = 0
8585
for name, _ in tensors.items():
86-
if name.endswith("k_proj") or name.endswith("v_proj"):
86+
if name.endswith("k_proj.weight") or name.endswith("v_proj.weight"):
8787
proj_linear_count += 1
8888
if name.endswith("k_proj.output_scale") or name.endswith("v_proj.output_scale"):
8989
output_scale_count += 1

0 commit comments

Comments
 (0)