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.
This commit is contained in:
Daniel Han 2026-04-16 20:31:49 +00:00
commit d323b47805

View file

@ -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
]