fix deepseek v2 input feat (#524)

casper-hansen · Jun 30, 2024 · 1768bac · 1768bac
1 parent 35d23db
commit 1768bac
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/awq/models/deepseek_v2.py b/awq/models/deepseek_v2.py
@@ -48,7 +48,7 @@ def get_layers_for_scaling(
                         module.self_attn.q_a_proj,
                         module.self_attn.kv_a_proj_with_mqa,
                     ],
-                    inp=input_feat["self_attn.q_proj"],
+                    inp=input_feat["self_attn.q_a_proj"],
                     module2inspect=module.self_attn,
                     kwargs=module_kwargs,
                 )