From d323b47805ed79831c450827d244b8a0330d2b21 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 16 Apr 2026 20:31:49 +0000
Subject: [PATCH 1/3] Add Gemma-4 to FORCE_FLOAT32 to prevent fp16 NaN in RL
 training

Gemma-4 produces NaN during GRPO training with dtype=torch.float16
due to torch.compile + fp16 autocast interaction in the backward pass.
Same class of issue as Gemma-3. Forces bfloat16 + float32 mixed
precision when users request fp16.
---
 unsloth/models/loader.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index fc91178d88..8e79749b0d 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -105,6 +105,8 @@ FORCE_FLOAT32 = [
     "gemma3,",  # Add comma bc gemma3 will match gemma3n
     "gemma3text",  # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3)
     "gemma3n",
+    "gemma4,",  # Add comma bc gemma4 will match gemma4_text
+    "gemma4text",  # Gemma4TextModel (standalone text-only Gemma4)
     "gpt_oss",
     "qwen3_5",  # Qwen3.5 GDN layers produce NaN grad norms in float16 training
 ]

From e5ee16b8a1462b65fc884d9dd707220039b00798 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 16 Apr 2026 21:55:00 +0000
Subject: [PATCH 2/3] Add FORCE_BFLOAT16 for Gemma-4 to prevent fp16 NaN in RL
 training

Gemma-4 produces NaN during GRPO training with dtype=torch.float16
due to torch.compile backward graph overflowing fp16 intermediates.
FORCE_FLOAT32 does not fix this because it still runs fp16 compute
(loads bf16 then casts to fp16 with no autocast). Instead, silently
switch to bfloat16 which has sufficient range to avoid overflow.

Adds a new FORCE_BFLOAT16 list separate from FORCE_FLOAT32. When a
user requests fp16 for Gemma-4 on a bf16-capable device, the dtype is
transparently switched to bfloat16 with a warning message. No changes
to the training pipeline or mixed precision behavior are needed.

Verified: 12 GRPO steps, ga=4, num_generations=2, all NaN-free.
---
 unsloth/models/loader.py | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 8e79749b0d..f92d348080 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -105,12 +105,17 @@ FORCE_FLOAT32 = [
     "gemma3,",  # Add comma bc gemma3 will match gemma3n
     "gemma3text",  # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3)
     "gemma3n",
-    "gemma4,",  # Add comma bc gemma4 will match gemma4_text
-    "gemma4text",  # Gemma4TextModel (standalone text-only Gemma4)
     "gpt_oss",
     "qwen3_5",  # Qwen3.5 GDN layers produce NaN grad norms in float16 training
 ]
 
+# Models that must use bfloat16 instead of float16.
+# torch.compile backward graphs overflow fp16 intermediates for these models.
+FORCE_BFLOAT16 = [
+    "gemma4,",  # Add comma bc gemma4 will match gemma4_text
+    "gemma4text",  # Gemma4TextModel (standalone text-only Gemma4)
+]
+
 global DISABLE_COMPILE_MODEL_NAMES
 # Must be alphabetically sorted for each entry
 
@@ -1375,6 +1380,20 @@ class FastModel(FastBaseModel):
                 os.environ["UNSLOTH_FORCE_FLOAT32"] = "1"
                 dtype = torch.bfloat16  # Change to bfloat16 loading
                 break
+        # Switch fp16 to bf16 for models whose torch.compile backward overflows fp16
+        global FORCE_BFLOAT16
+        for disable_name in FORCE_BFLOAT16:
+            if (
+                disable_name.lower()
+                == model_type_arch.lower().replace("-", "").replace("_", "")
+                or disable_name.lower() in model_types_all
+            ) and dtype == torch.float16 and SUPPORTS_BFLOAT16:
+                logger.warning_once(
+                    f"Unsloth: {model_type_arch} does not support float16 training. "
+                    f"Switching to bfloat16."
+                )
+                dtype = torch.bfloat16
+                break
         # Apply gradient checkpointing with smart heuristics
         use_gradient_checkpointing = apply_unsloth_gradient_checkpointing(
             use_gradient_checkpointing, max_seq_length, dtype

From e264d42413c153b7ce3cf4eec55dccfc79b2deb9 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 16 Apr 2026 21:55:14 +0000
Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 unsloth/models/loader.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index f92d348080..a42b7f38a7 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -1384,10 +1384,14 @@ class FastModel(FastBaseModel):
         global FORCE_BFLOAT16
         for disable_name in FORCE_BFLOAT16:
             if (
-                disable_name.lower()
-                == model_type_arch.lower().replace("-", "").replace("_", "")
-                or disable_name.lower() in model_types_all
-            ) and dtype == torch.float16 and SUPPORTS_BFLOAT16:
+                (
+                    disable_name.lower()
+                    == model_type_arch.lower().replace("-", "").replace("_", "")
+                    or disable_name.lower() in model_types_all
+                )
+                and dtype == torch.float16
+                and SUPPORTS_BFLOAT16
+            ):
                 logger.warning_once(
                     f"Unsloth: {model_type_arch} does not support float16 training. "
                     f"Switching to bfloat16."