granite force layernorm upcast (#2799)

This commit is contained in:
DoubleMathew 2025-06-28 07:27:56 -05:00 committed by GitHub
commit ebe935ee74

View file

@ -556,6 +556,10 @@ class FastModel(FastBaseModel):
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)"
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"
os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1"
elif "olmo-2" in lowered_model_name and transformers_version < Version("4.50.0.dev0"):
raise RuntimeError("Unsloth: OLMo-2 only works on transformers >= 4.50.0." + NIGHTLY)
elif "gemma-3n" in lowered_model_name: