Update gemma2.py

unslothai · Jun 30, 2024 · 73148cd · 73148cd
1 parent f2777ed
commit 73148cd
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py
@@ -279,10 +279,7 @@ def Gemma2DecoderLayer_forward(
             attention_mask = attention_mask[:, :, :, -self.sliding_window :]
 
     residual = hidden_states
-
-    hidden_states = self.input_layernorm(hidden_states)
-
-    # Self Attention
+    hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states, gemma = True)
     hidden_states, self_attn_weights, present_key_value = self.self_attn(
         hidden_states=hidden_states,
         attention_mask=attention_mask,
@@ -292,13 +289,13 @@ def Gemma2DecoderLayer_forward(
         use_cache=use_cache,
         cache_position=cache_position,
     )
-    hidden_states = self.post_attention_layernorm(hidden_states)
+    hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, gemma = True)
     hidden_states = residual + hidden_states
 
     residual = hidden_states
-    hidden_states = self.pre_feedforward_layernorm(hidden_states)
+    hidden_states = fast_rms_layernorm(self. pre_feedforward_layernorm, hidden_states, gemma = True)
     hidden_states = self.mlp(hidden_states)
-    hidden_states = self.post_feedforward_layernorm(hidden_states)
+    hidden_states = fast_rms_layernorm(self.post_feedforward_layernorm, hidden_states, gemma = True)
     hidden_states = residual + hidden_states
 
     outputs = (hidden_states,)