From 67e3251577bb08a69b075deba09ba73121fb54a9 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 7 Jan 2025 00:50:56 -0800
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 824a3ccf15..8a1d2c99b2 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -498,7 +498,9 @@ def LlamaDecoderLayer_fast_forward(
         hidden_states += residual
     else:
         residual = hidden_states
+        print(501, hidden_states.dtype)
         hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states)
+        print(503, hidden_states.dtype)
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states       = hidden_states,
             causal_mask         = causal_mask,
@@ -510,12 +512,16 @@ def LlamaDecoderLayer_fast_forward(
             padding_mask        = padding_mask,
             position_embeddings = position_embeddings,
         )
+        print(515, hidden_states.dtype)
         hidden_states = residual + hidden_states
 
         # Fully Connected
         residual = hidden_states
+        print(520, hidden_states.dtype)
         hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states)
+        print(522, hidden_states.dtype)
         hidden_states = self.mlp(hidden_states)
+        print(524, hidden_states.dtype)
         hidden_states = residual + hidden_states
     pass