From e5ee16b8a1462b65fc884d9dd707220039b00798 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 16 Apr 2026 21:55:00 +0000 Subject: [PATCH] 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