Fix starcoder2 fused norm (#442)

casper-hansen · Apr 12, 2024 · 0e52a5c · 0e52a5c
1 parent e9f6269
commit 0e52a5c
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/awq/models/starcoder2.py b/awq/models/starcoder2.py
@@ -110,13 +110,9 @@ def fuse_transformer(self):
                 module.self_attn.k_proj,
                 module.self_attn.v_proj,
             )
-            norm_1 = FasterTransformerRMSNorm(
-                module.input_layernorm.weight, module.input_layernorm.eps
-            )
-            norm_2 = FasterTransformerRMSNorm(
-                module.post_attention_layernorm.weight,
-                module.post_attention_layernorm.eps,
-            )
+            # SC2 use normal LayerNorm
+            norm_1 = module.input_layernorm
+            norm_2 = module.post_attention_layernorm
             blocks.append(
                 LlamaLikeBlock(
                     hidden_size=self.model.config.hidden_size,