Update loader.py

unsloth/models/loader.py

@@ -59,7 +59,15 @@ if SUPPORTS_GEMMA2:
     from .gemma2 import FastGemma2Model
 pass
 import torch
-
+from ._utils import (
+    patch_compiling_bitsandbytes,
+    patch_model_and_tokenizer,
+    prepare_model_for_kbit_training,
+    patch_unsloth_smart_gradient_checkpointing,
+    patch_compiled_autograd,
+    process_vision_info,
+    unsloth_compile_transformers,
+)
 
 class FastLanguageModel(FastLlamaModel):
     @staticmethod
@@ -87,6 +95,10 @@ class FastLanguageModel(FastLlamaModel):
         *args, **kwargs,
     ):
         if token is None: token = get_token()
+        assert (dtype is None or dtype == torch.float16 or dtype == torch.bfloat16)
+
+        if use_gradient_checkpointing == "unsloth":
+            patch_unsloth_smart_gradient_checkpointing(dtype = dtype)
 
         if fast_inference:
             if importlib.util.find_spec("vllm") is None:
@@ -367,15 +379,6 @@ class FastLanguageModel(FastLlamaModel):
 pass
 
 
-from ._utils import (
-    patch_compiling_bitsandbytes,
-    patch_model_and_tokenizer,
-    prepare_model_for_kbit_training,
-    patch_unsloth_smart_gradient_checkpointing,
-    patch_compiled_autograd,
-    process_vision_info,
-    unsloth_compile_transformers,
-)
 from ..kernels import (
     patch_loss_functions,
     post_patch_loss_function,
@@ -404,6 +407,7 @@ class FastVisionModel(FastBaseVisionModel):
         *args, **kwargs,
     ):
         if token is None: token = get_token()
+        assert (dtype is None or dtype == torch.float16 or dtype == torch.bfloat16)
 
         patch_compiled_autograd()
         patch_compiling_bitsandbytes()