diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 7800ec2b0a..f7e08ed10d 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -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: