From acfe670357f27ad5625c4e1804f0d4e47120c675 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 10 Feb 2026 01:40:13 -0800 Subject: [PATCH] Fix EmbeddingGemma float16 NaN via FORCE_FLOAT32 for gemma3_text (#4014) * Fix EmbeddingGemma float16 NaN by adding gemma3_text to FORCE_FLOAT32 and SDPA lists * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Hanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/loader.py | 2 ++ unsloth/models/sentence_transformer.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 97add13f2d..9d2d3d9b06 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -100,6 +100,7 @@ global FORCE_FLOAT32 # Forces float32 precision since float16 goes to infinity FORCE_FLOAT32 = [ "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3) "gemma3n", "gpt_oss", ] @@ -116,6 +117,7 @@ global DISABLE_SDPA_MODEL_NAMES # Disables some SDPA modules since it's wrong DISABLE_SDPA_MODEL_NAMES = [ "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore ] diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index a3ea950420..ad59165a50 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -2089,6 +2089,20 @@ def _patch_sentence_transformer_trainer(): # Call original __init__ _original_init(self, *args, **kwargs) + # Disable mixed precision when FORCE_FLOAT32 is active (matches rl.py behavior) + if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": + if hasattr(self, "args") and self.args is not None: + if self.args.fp16 or self.args.bf16: + print( + "Unsloth: Switching to float32 training since model cannot work with float16" + ) + self.args.fp16 = False + self.args.bf16 = False + if hasattr(self.args, "bf16_full_eval"): + self.args.bf16_full_eval = False + if hasattr(self.args, "fp16_full_eval"): + self.args.fp16_full_eval = False + SentenceTransformerTrainer.__init__ = _patched_init SentenceTransformerTrainer._unsloth_auto_compile_patched = True