ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix EmbeddingGemma float16 NaN via FORCE_FLOAT32 for gemma3_text (#4014)

Browse source 
* 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 <danielhanchen@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-02-10 01:40:13 -08:00 • committed by GitHub
commit acfe670357
2 changed files with 16 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/models/loader.py
View file

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

14
unsloth/models/sentence_transformer.py
View file

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