From 27cd2b3a0de189b399bb2021bd77edd469b2bca3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 5 Nov 2024 14:40:57 -0800
Subject: [PATCH] CE Loss

---
 unsloth/kernels/cross_entropy_loss.py | 10 ++++++----
 unsloth/models/_utils.py              |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index f0193c74d8..13d90baafe 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -181,10 +181,10 @@ def _chunked_cross_entropy_forward(
 pass
 
 
-@triton.heuristics({
-    "DO_SOFTCAPPING":   lambda args: bool(args["DO_SOFTCAPPING"  ]),
-    "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
-})
+# @triton.heuristics({
+#     "DO_SOFTCAPPING":   lambda args: bool(args["DO_SOFTCAPPING"  ]),
+#     "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
+# })
 @triton.jit
 def _cross_entropy_backward(
     logits_ptr        ,
@@ -345,6 +345,8 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
         n_rows, vocab_size = logits.shape
 
         BLOCK_SIZE : int = 4096
+        div : int
+        mod : int
         div, mod = divmod(vocab_size, BLOCK_SIZE)
         n_blocks : int = div + (mod != 0)
 
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index b5c17434c7..bb004d588d 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -393,7 +393,7 @@ UNSLOTH_COMPILE_DEBUG   = os.environ.get("UNSLOTH_COMPILE_DEBUG",   "0") == "1"
 UNSLOTH_COMPILE_MAXIMUM = os.environ.get("UNSLOTH_COMPILE_MAXIMUM", "0") == "1"
 patch_torch_compile(
     debug = UNSLOTH_COMPILE_DEBUG,
-    O3 = UNSLOTH_COMPILE_MAXIMUM,
+    O3    = UNSLOTH_COMPILE_MAXIMUM,
 )
 
 torch_compile_options = {