From e517fd384da39edccf6979edefd9b07054c009bb Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 30 Oct 2024 13:11:43 -0700
Subject: [PATCH] Bug fixes

---
 unsloth/__init__.py                   | 14 +++++++-------
 unsloth/kernels/cross_entropy_loss.py | 10 +++++++---
 unsloth/models/_utils.py              |  6 +++++-
 unsloth/models/llama.py               |  1 +
 4 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index 458c2696bc..109e1c6d2f 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -27,13 +27,6 @@ import numpy as np
 #     pass
 # pass
 
-# Check for unsloth_zoo
-try:
-    import unsloth_zoo
-except:
-    raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth-zoo`")
-pass
-
 # Unsloth currently does not work on multi GPU setups - sadly we are a 2 brother team so
 # enabling it will require much more work, so we have to prioritize. Please understand!
 # We do have a beta version, which you can contact us about!
@@ -165,6 +158,13 @@ if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ
     pass
 pass
 
+# Check for unsloth_zoo
+try:
+    import unsloth_zoo
+except:
+    raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth-zoo`")
+pass
+
 from .models import *
 from .save import *
 from .chat_templates import *
diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index debd037b64..a2337a14d9 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -17,7 +17,7 @@ import triton.language as tl
 import torch
 from .utils import calculate_settings, MAX_FUSED_SIZE, triton_tanh
 from transformers.models.llama.modeling_llama import logger
-
+from packaging.version import Version
 
 @triton.heuristics({
     "DO_SOFTCAPPING":   lambda args: args["DO_SOFTCAPPING"  ],
@@ -352,7 +352,6 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
 pass
 
 
-# @torch._disable_dynamo
 def fast_cross_entropy_loss(
     logits,
     labels,
@@ -380,6 +379,9 @@ def fast_cross_entropy_loss(
         n_items = torch.count_nonzero(labels != -100)
     return loss.sum() / n_items
 pass
+if Version(torch.__version__) < Version("2.5.0"):
+    fast_cross_entropy_loss = torch._disable_dynamo(fast_cross_entropy_loss)
+pass
 
 
 from transformers.models.llama.modeling_llama import (
@@ -475,7 +477,6 @@ def unpatch_llama_for_causal_lm():
 pass
 
 
-# @torch._disable_dynamo
 def UnslothForCausalLMLoss(
     logits, labels, vocab_size: int, num_items_in_batch: int = None, ignore_index: int = -100, **kwargs
 ):
@@ -490,6 +491,9 @@ def UnslothForCausalLMLoss(
     )
     return loss
 pass
+if Version(torch.__version__) < Version("2.5.0"):
+    UnslothForCausalLMLoss = torch._disable_dynamo(UnslothForCausalLMLoss)
+pass
 
 
 def patch_transformers_losses():
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 97b2ea7b7b..a39bc58db0 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -760,7 +760,9 @@ def get_statistics():
     # We log some basic stats about which environment is being used.
     # We simply download a README.md file from HF - all data is made public.
     # This is simply so we can check if some envs are broken or not.
-    # You can disable this by commenting the below out
+    # You can disable this by setting UNSLOTH_DISABLE_STATISTICS
+    import os
+    if "UNSLOTH_DISABLE_STATISTICS" in os.environ: return
     from huggingface_hub.utils import disable_progress_bars, enable_progress_bars, are_progress_bars_disabled
     disabled = False
     if not are_progress_bars_disabled():
@@ -1295,6 +1297,7 @@ def patch_gradient_accumulation_fix(Trainer):
     # Fixes gradient accumulation 
     import inspect
     if hasattr(Trainer, "get_batch_samples"):
+        if Trainer.get_batch_samples.__name__ == "_unsloth_get_batch_samples": return
         if \
             not inspect.getsource(Trainer.get_batch_samples).strip()\
             .endswith("return batch_samples, num_items_in_batch"):
@@ -1321,6 +1324,7 @@ def patch_gradient_accumulation_fix(Trainer):
     pass
 
     # Also fix up loss scaling ie negate loss *= self.args.gradient_accumulation_steps
+    if Trainer.training_step.__name__ == "_unsloth_training_step": return
     if "num_items_in_batch" not in inspect.signature(Trainer.training_step).parameters: return
 
     function = inspect.getsource(Trainer.training_step)
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 65e2d773e8..c0175bbfaf 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1518,6 +1518,7 @@ class FastLlamaModel:
         pass
         # Return old flag
         os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer
+        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
         model_patcher.pre_patch()
         get_statistics() # For debugging - we use a download counter to see if environments are not breaking