Fix starcoder2 fused norm (#442)

0e52a5c2 · TechxGenus · GitHub · e9f62694 · 0e52a5c2
Unverified Commit 0e52a5c2 authored Apr 13, 2024 by TechxGenus Committed by GitHub Apr 12, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 7 deletions

awq/models/starcoder2.py awq/models/starcoder2.py +3 -7

No files found.
--- a/awq/models/starcoder2.py
+++ b/awq/models/starcoder2.py
@@ -110,13 +110,9 @@ class Starcoder2Fuser:
                module.self_attn.k_proj,
                module.self_attn.v_proj,
            )
-            norm_1 = FasterTransformerRMSNorm(
+            # SC2 use normal LayerNorm
-                module.input_layernorm.weight, module.input_layernorm.eps
+            norm_1 = module.input_layernorm
-            )
+            norm_2 = module.post_attention_layernorm
-            norm_2 = FasterTransformerRMSNorm(
-                module.post_attention_layernorm.weight,
-                module.post_attention_layernorm.eps,
-            )
            blocks.append(
                LlamaLikeBlock(
                    hidden_size=self.model.config.hidden_size,