Gemma 3N fixes

This commit is contained in:
Daniel Han 2025-06-30 06:51:49 -07:00
commit e1654ec5a3
2 changed files with 13 additions and 4 deletions

View file

@ -562,7 +562,9 @@ class FastModel(FastBaseModel):
# Sesame
elif "csm-1b" in lowered_model_name:
os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" # Sesame fails
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "all;torch.float32;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);"
# Granite 4
elif 'granite-4' in lowered_model_name:
# granite-4 rms norms are stored as 16 bit, but we upcast
@ -574,7 +576,11 @@ class FastModel(FastBaseModel):
# Gemma 3N
elif "gemma-3n" in lowered_model_name:
os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1"
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "float16;torch.float16;torch.float16;if name.endswith(('.conv')): module.to(torch.float32)"
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \
"float16;torch.float16;torch.float16;"\
"if name.endswith(('.conv')): module.to(torch.float32);"\
"from unsloth_zoo.temporary_patches.gemma3n import patch_Gemma3nConvNormAct_forward; patch_Gemma3nConvNormAct_forward()"
if transformers_version < Version("4.53.0"):
raise RuntimeError("Unsloth: Gemma 3N only works on transformers >= 4.53.0" + LATEST)
else:

View file

@ -350,8 +350,8 @@ 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(";") == 3
checker, _dtype, _bnb_compute_dtype, _custom_datatype = custom_datatype.split(";", 3)
assert custom_datatype.count(";") >= 4
checker, _dtype, _bnb_compute_dtype, _custom_datatype, execute_code = custom_datatype.split(";", 4)
# Allow custom dtypes on all runs
allow_all_runs = (checker == "all")
@ -363,6 +363,9 @@ class FastBaseModel:
bnb_compute_dtype = eval(_bnb_compute_dtype)
correct_dtype = bnb_compute_dtype
custom_datatype = _custom_datatype
# Execute code as well
print(execute_code)
exec(execute_code)
pass
pass