From aa96d65534285dfa86f00bea64370c80e893f79b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 7 Aug 2025 09:38:24 -0700 Subject: [PATCH] Update loader.py --- unsloth/models/loader.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 94a0c526db..bc98c0dc19 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -592,12 +592,14 @@ class FastModel(FastBaseModel): "os.environ['TRITON_F32_DEFAULT'] = 'ieee';" elif "gpt-oss" in lowered_model_name: os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" - os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \ - "all;None;None;"\ - "x = 'gate_up_proj_bias'\n"\ - "if hasattr(module, x): setattr(module, x, torch.nn.Parameter(getattr(module, x).to(torch.float32)) if isinstance(getattr(module, x), torch.nn.Parameter) else getattr(module, x).to(torch.float32))\n"\ - "x = 'down_proj_bias'\n"\ - "if hasattr(module, x): setattr(module, x, torch.nn.Parameter(getattr(module, x).to(torch.float32)) if isinstance(getattr(module, x), torch.nn.Parameter) else getattr(module, x).to(torch.float32))\n;" + if not model_name: + # Only upcast MoE biases for MXFP4, not BnB + os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \ + "all;None;None;"\ + "x = 'gate_up_proj_bias'\n"\ + "if hasattr(module, x): setattr(module, x, torch.nn.Parameter(getattr(module, x).to(torch.float32)) if isinstance(getattr(module, x), torch.nn.Parameter) else getattr(module, x).to(torch.float32))\n"\ + "x = 'down_proj_bias'\n"\ + "if hasattr(module, x): setattr(module, x, torch.nn.Parameter(getattr(module, x).to(torch.float32)) if isinstance(getattr(module, x), torch.nn.Parameter) else getattr(module, x).to(torch.float32))\n;" else: for check_model_name in DISABLE_COMPILE_MODEL_NAMES: if check_model_name in lowered_model_name: