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

Gemma 3N fixes

Browse source
This commit is contained in:
Daniel Han 2025-06-30 06:01:50 -07:00
commit d0b03f80b2
2 changed files with 16 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/models/loader.py
View file

@ -557,7 +557,7 @@ class FastModel(FastBaseModel):
raise RuntimeError("Unsloth: Cohere's Command model only works on transformers >= 4.50.0." + NIGHTLY)
elif "csm-1b" in lowered_model_name:
os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" # Sesame fails
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "torch.float16;if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16)"
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "all;torch.float32;torch.float16;if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16)"
elif 'granite-4' in lowered_model_name:
# granite-4 rms norms are stored as 16 bit, but we upcast
os.environ["UNSLOTH_UPCAST_LAYERNORM"] = "1"
@ -566,6 +566,7 @@ class FastModel(FastBaseModel):
raise RuntimeError("Unsloth: OLMo-2 only works on transformers >= 4.50.0." + NIGHTLY)
elif "gemma-3n" in lowered_model_name:
os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1"
s.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "float16;torch.float16;torch.float16;if name.endswith(('.conv')): module.to(torch.float32)"
if transformers_version < Version("4.53.0"):
raise RuntimeError("Unsloth: Gemma 3N only works on transformers >= 4.53.0" + LATEST)
else:

19
unsloth/models/vision.py
View file

@ -352,11 +352,20 @@ class FastBaseModel:
correct_dtype = None
if os.environ.get("UNSLOTH_FORCE_CUSTOM_DTYPE", "") != "":
custom_datatype = os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"]
assert custom_datatype.count(";") == 1
bnb_compute_dtype, custom_datatype = custom_datatype.split(";", 1)
dtype = torch.float32
bnb_compute_dtype = eval(bnb_compute_dtype)
correct_dtype = bnb_compute_dtype
assert custom_datatype.count(";") == 3
checker, _dtype, _bnb_compute_dtype, _custom_datatype = custom_datatype.split(";", 3)
# Allow custom dtypes on all runs
allow_all_runs = (checker == "all")
# Allow only on float16 datatypes
allow_float16_runs = (checker == "float16" and dtype == torch.float16)
if allow_all_runs or allow_float16_runs:
dtype = eval(_dtype)
bnb_compute_dtype = eval(_bnb_compute_dtype)
correct_dtype = bnb_compute_dtype
custom_datatype = _custom_datatype
pass
pass
# Stop SDPA for some archs like Pixtral / Mistral3