From 762185dd7b10595d400d640f26a002b81e8f2c76 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 2 Jan 2025 18:25:25 -0800
Subject: [PATCH] requires grad

---
 unsloth/__init__.py      | 21 ++++++++++-----------
 unsloth/models/_utils.py |  6 ------
 unsloth/models/vision.py |  3 +++
 3 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index 90d2a63519..bbeded9fc6 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -17,16 +17,6 @@ from packaging.version import Version
 import os, re, subprocess, inspect
 import numpy as np
 
-# # Define a list of modules to check
-# MODULES_TO_CHECK = ["bitsandbytes"]
-
-# # Check if any of the modules in the list have been imported
-# for module in MODULES_TO_CHECK:
-#     if module in sys.modules:
-#         raise ImportError(f"Unsloth: Please import Unsloth before {module}.")
-#     pass
-# 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!
@@ -201,9 +191,18 @@ pass
 
 # Check for unsloth_zoo
 try:
+    unsloth_zoo_version = importlib_version("unsloth_zoo")
+    if Version(unsloth_zoo_version) < Version("2025.1.1"):
+        try:
+            os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo")
+        except:
+            try:
+                os.system("pip install --upgrade --no-cache-dir --no-deps --user unsloth_zoo")
+            except:
+                raise ImportError("Unsloth: Please update unsloth_zoo via `pip install --upgrade --no-cache-dir --no-deps unsloth_zoo`")
     import unsloth_zoo
 except:
-    raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth-zoo`")
+    raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth_zoo`")
 pass
 
 from .models import *
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 33fb36e8b6..098f5c3e47 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -110,9 +110,6 @@ from unsloth_zoo.compiler import (
     get_transformers_model_type,
     unsloth_compile_transformers as _unsloth_compile_transformers,
 )
-from unsloth_zoo.peft_utils import (
-    requires_grad_for_gradient_checkpointing,
-)
 
 # =============================================
 # Disable some warnings which can get annoying
@@ -559,9 +556,6 @@ def prepare_model_for_kbit_training(
             def make_inputs_require_grad(module, input, output):
                 output.requires_grad_(True)
             model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
-
-        # Enable grads on non language models as well
-        requires_grad_for_gradient_checkpointing(model)
     pass
 
     return model
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 2dc4b88dfa..51450aa0d9 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -30,6 +30,7 @@ from transformers import set_seed as transformers_set_seed
 from unsloth_zoo.peft_utils import (
     get_peft_regex,
     SKIP_QUANTIZATION_MODULES,
+    requires_grad_for_gradient_checkpointing,
 )
 from triton import __version__ as triton_version
 
@@ -275,6 +276,8 @@ class FastBaseVisionModel:
             use_gradient_checkpointing = use_gradient_checkpointing,
         )
         model = get_peft_model(model, lora_config)
+        # Enable gradients on modules which are trainable
+        requires_grad_for_gradient_checkpointing(model)
 
         model = FastBaseVisionModel.patch_peft_model(model, use_gradient_checkpointing)