Bug fixes

pyproject.toml

@@ -33,10 +33,10 @@ exclude = ["images*"]
 
 [project.optional-dependencies]
 huggingface = [
-    "unsloth_zoo",
+    "unsloth_zoo>=2024.11.1",
     "packaging",
     "tyro",
-    "transformers>=4.44.2",
+    "transformers>=4.46.1",
     "datasets>=2.16.0",
     "sentencepiece>=0.2.0",
     "tqdm",
@@ -244,10 +244,10 @@ colab-ampere-torch220 = [
     "flash-attn>=2.6.3",
 ]
 colab-new = [
-    "unsloth_zoo",
+    "unsloth_zoo>=2024.11.1",
     "packaging",
     "tyro",
-    "transformers>=4.44.2",
+    "transformers>=4.46.1",
     "datasets>=2.16.0",
     "sentencepiece>=0.2.0",
     "tqdm",

unsloth/__init__.py

@@ -27,13 +27,6 @@
 #     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!
@@ -60,6 +53,14 @@
 # Reduce VRAM usage by reducing fragmentation
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
+# Hugging Face Hub faster downloads
+if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+pass
+
+# Log Unsloth is being used
+os.environ["UNSLOTH_IS_PRESENT"] = "1"
+
 try:
     import torch
 except ModuleNotFoundError:
@@ -71,12 +72,6 @@
     raise exception
 pass
 
-# Hugging Face Hub faster downloads (only enable during Colab and Kaggle sessions)
-keynames = "\n" + "\n".join(os.environ.keys())
-if "\nCOLAB_"  in keynames or "\nKAGGLE_" in keynames:
-    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-pass
-
 # We support Pytorch 2
 # Fixes https://github.com/unslothai/unsloth/issues/38
 torch_version = torch.__version__.split(".")
@@ -165,6 +160,13 @@ def is_bf16_supported(): return SUPPORTS_BFLOAT16
     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 *

unsloth/kernels/__init__.py

@@ -14,8 +14,8 @@
 
 from .cross_entropy_loss import (
     fast_cross_entropy_loss,
-    patch_llama_for_causal_lm,
-    unpatch_llama_for_causal_lm,
+    post_patch_loss_function,
+    patch_loss_functions,
 )
 from .rms_layernorm import (
     fast_rms_layernorm,
@@ -25,7 +25,6 @@
 from .layernorm import (
     fast_layernorm,
     patch_layernorm,
-    unpatch_layernorm,
 )
 from .rope_embedding import fast_rope_embedding, inplace_rope_embedding
 from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
@@ -54,8 +53,12 @@
     create_flex_attention_sliding_window_mask,
 )
 
-try:
-    print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.")
-except:
-    print("Unsloth: Will patch your computer to enable 2x faster free finetuning.")
+import os
+if "UNSLOTH_ZOO_IS_PRESENT" not in os.environ:
+    try:
+        print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.")
+    except:
+        print("Unsloth: Will patch your computer to enable 2x faster free finetuning.")
+    pass
 pass
+del os