From 9a356a7f7945551dafd284f8a2382aed1a6fd8b1 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 17 Mar 2025 04:41:44 -0700
Subject: [PATCH] move forced float32

---
 unsloth/models/_utils.py | 20 ++++++++++++++++++++
 unsloth/models/loader.py | 10 +++++++++-
 unsloth/models/vision.py | 20 +++-----------------
 3 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 2375fff4d8..f84d80f280 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -121,6 +121,11 @@ from unsloth_zoo.temporary_patches import (
 for temporary_patch in TEMPORARY_PATCHES:
     temporary_patch()
 
+global FORCE_FLOAT32
+FORCE_FLOAT32 = [
+    "gemma3",
+]
+
 # =============================================
 # Disable some warnings which can get annoying
 warnings.filterwarnings(action = "ignore", category = UserWarning,    module = "torch")
@@ -1127,6 +1132,7 @@ pass
 
 
 def unsloth_compile_transformers(
+    dtype,
     model_name,
     token                   = None,
     revision                = None,
@@ -1176,6 +1182,20 @@ def unsloth_compile_transformers(
 
     if disable: return
 
+    # Set forced float32 env flag
+    os.environ["UNSLOTH_FORCE_FLOAT32"] = "0"
+    do_forced_float32 = False
+    for disable_name in FORCE_FLOAT32:
+        if (disable_name.lower() == model_types[1].lower() or \
+            disable_name.lower() in model_name.lower()) and \
+            dtype == torch.float16:
+
+            print(f"Unsloth: Using float16 precision for {model_type_arch} won't work! Using float32.")
+            os.environ["UNSLOTH_FORCE_FLOAT32"] = "1"
+            do_forced_float32 = True
+            break
+    pass
+
     for model_type in model_types:
         _unsloth_compile_transformers(
             model_type,
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 262d403b3a..f73f0d3ec6 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -460,7 +460,14 @@ class FastModel(FastBaseModel):
         *args, **kwargs,
     ):
         if token is None: token = get_token()
-        assert (dtype is None or dtype in (torch.float16, torch.bfloat16, torch.float32))
+
+        SUPPORTS_BFLOAT16 = is_bfloat16_supported()
+        if dtype is None:
+            dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
+        elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
+            logger.warning_once("Device does not support bfloat16. Will change to float16.")
+            dtype = torch.float16
+        assert(dtype in (torch.float16, torch.bfloat16, torch.float32))
 
         patch_compiled_autograd()
         patch_compiling_bitsandbytes()
@@ -614,6 +621,7 @@ class FastModel(FastBaseModel):
         with redirector:
             patch_loss_functions(torch_compile = False)
             model_types = unsloth_compile_transformers(
+                dtype                   = dtype,
                 model_name              = model_name,
                 sdpa_dynamic_mask       = True,
                 sdpa_bool_masks         = True,
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 53ead28ac7..50df3999ff 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -65,11 +65,6 @@ __all__ = [
     "FastBaseModel",
 ]
 
-global FORCE_FLOAT32
-FORCE_FLOAT32 = [
-    "gemma3",
-]
-
 global FORCE_EAGER_ATTENTION
 FORCE_EAGER_ATTENTION = [
     "pixtral",    # Pixtral SDPA not implemented
@@ -215,20 +210,11 @@ class FastBaseModel:
 
         assert(dtype in (torch.float16, torch.bfloat16, torch.float32))
 
-        global FORCE_FLOAT32
-        os.environ["UNSLOTH_FORCE_FLOAT32"] = "0"
         bnb_compute_dtype = dtype
         do_forced_float32 = False
-        for disable_name in FORCE_FLOAT32:
-            if (disable_name.lower() == model_type_arch.lower() or \
-                disable_name.lower() in model_name.lower()) and \
-                dtype == torch.float16:
-
-                print(f"Unsloth: Using float16 precision for {model_type_arch} won't work! Using float32.")
-                os.environ["UNSLOTH_FORCE_FLOAT32"] = "1"
-                bnb_compute_dtype = torch.float16
-                do_forced_float32 = True
-                break
+        if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
+            bnb_compute_dtype = torch.float16
+            do_forced_float32 = True
         pass
 
         global FORCE_EAGER_ATTENTION