From 100f2a3242fcd9c4185cc6944427ccb82aab1899 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 13 Feb 2025 17:05:04 -0800
Subject: [PATCH] Update llama.py

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

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 817b014ac0..188c12ba96 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1068,6 +1068,7 @@ def CausalLM_fast_forward(fast_forward_inference):
             print(1068)
         pass
         hidden_states = outputs[0]
+        print(1071)
 
         bsz, q_len, hd = hidden_states.shape
         lm_head = self.lm_head.weight
@@ -1084,6 +1085,8 @@ def CausalLM_fast_forward(fast_forward_inference):
             RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
             # < 1024 Normal Unsloth uses less VRAM!
             if bsz*q_len <= 1024: RETURN_LOGITS = True
+
+            print(1089)
             
             if not RETURN_LOGITS and HAS_CUT_CROSS_ENTROPY and labels is not None:
 
@@ -1095,6 +1098,8 @@ def CausalLM_fast_forward(fast_forward_inference):
                     num_items_in_batch = n_items,
                     logit_softcapping  = logit_softcapping,
                 )
+
+                print(1102, loss)
                 if not return_dict:
                     output = (logits,) + outputs[1:]
                     return (loss,) + output if loss is not None else output
@@ -1108,6 +1113,7 @@ def CausalLM_fast_forward(fast_forward_inference):
                 )
                 return output
             pass
+            print(1116, hidden_states.dtype, hidden_states.shape)
             logits = self.lm_head(hidden_states.to(dtype))
         pass
 
@@ -1117,6 +1123,7 @@ def CausalLM_fast_forward(fast_forward_inference):
         else:
             raise TypeError("Unsloth: torch_dtype for models is not bfloat16, float16 or float32!")
         pass
+        print(1126)
 
         loss = None
         logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
@@ -1142,6 +1149,7 @@ def CausalLM_fast_forward(fast_forward_inference):
                 logit_scaling     = logit_scaling,
                 n_items           = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None),
             )
+            print(1152, loss)
         else:
             if logit_scaling != 0:
                 if logits.requires_grad:
@@ -1166,7 +1174,7 @@ def CausalLM_fast_forward(fast_forward_inference):
         if not return_dict:
             output = (logits,) + outputs[1:]
             return (loss,) + output if loss is not None else output
-
+        print(1177, loss, logits.shape, logits.dtype)
         return CausalLMOutputWithPast(
             loss=loss,
             logits=logits,