From 86a2725cb694c8f4bb3c6433748864cc7fa7708f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 14 Aug 2025 05:06:58 -0700
Subject: [PATCH] Fix GPT OSS (#3154)

* Fix mamba

* Update loader.py

* Update vision.py

* Update loader.py

* Filter vLLM standby logs (#3131)

* filter vLLM standby logs

* safeguard standby logger patch

* Update unsloth/models/_utils.py

* Update unsloth/models/_utils.py

* Update unsloth/models/_utils.py

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Update loader.py

* Add scaler

* Update llama.py

* Update _utils.py

* Versioning

* GPT OSS fix

* GPT OSS fix

* Update loader.py

* Update vision.py

* Update vision.py

* Update loader.py

---------

Co-authored-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
---
 unsloth/models/_utils.py |  2 ++
 unsloth/models/loader.py | 13 +++++++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index d1df57ad5c..ab2694fde1 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -68,6 +68,7 @@ __all__ = [
     "patch_fast_lora",
     "validate_loftq_config",
     "RaiseUninitialized",
+    "dequantize_module_weight",
 ]
 
 import torch
@@ -724,6 +725,7 @@ pass
 # Weirdly LoraLayer.update_layer downcasts PEFT layers to float16??
 # For mixed precision, we need it to be in float32 not float16.
 from peft import __version__ as peft_version
+from peft.utils.integrations import dequantize_module_weight
 if Version(peft_version) < Version("0.12.0"):
     from peft.tuners.lora.layer import LoraLayer
     try:
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 7ac27158a2..d0b7d4dc4c 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -618,9 +618,18 @@ class FastModel(FastBaseModel):
                 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"\
+                    "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:
+                # Set down projection compute dtype to be float32 for float16 machines
+                os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = \
+                    "all;None;None;"\
+                    "if 'down_projs' in name and hasattr(module, 'compute_dtype') and "\
+                    "torch.amax(dequantize_module_weight(module)) >= 1024:"\
+                    "module._pre_set_compute_dtype = torch.float32\n"\
+                    ";"
         else:
             for check_model_name in DISABLE_COMPILE_MODEL_NAMES:
                 if check_model_name in lowered_model_name: