From e3af5dc95dcfa116c4f266cb45167761d1dd4b14 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 7 Aug 2025 02:57:19 -0700
Subject: [PATCH] GPT OSS fixes

---
 unsloth/models/_utils.py |  3 ++-
 unsloth/models/loader.py | 20 ++++++--------------
 unsloth/models/vision.py |  6 ++++--
 3 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 2e88e213e6..a036970b40 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -76,6 +76,7 @@ platform_system = platform_system()
 import numpy as np
 import contextlib
 import re
+import functools
 import warnings, subprocess, re, inspect, psutil, os, math
 from unsloth_zoo.utils import Version
 from unsloth import DEVICE_TYPE, DEVICE_COUNT
@@ -422,6 +423,7 @@ HAS_FLASH_ATTENTION_SOFTCAPPING = False
 
 if DEVICE_TYPE == "cuda":
     major_version, minor_version = torch.cuda.get_device_capability()
+    torch.cuda.get_device_capability = functools.cache(torch.cuda.get_device_capability)
 
     if major_version >= 8:
         SUPPORTS_BFLOAT16 = True
@@ -586,7 +588,6 @@ UNSLOTH_COMPILE_DEBUG         = os.environ.get("UNSLOTH_COMPILE_DEBUG",
 UNSLOTH_COMPILE_MAXIMUM       = os.environ.get("UNSLOTH_COMPILE_MAXIMUM",       "0") == "1"
 UNSLOTH_COMPILE_IGNORE_ERRORS = os.environ.get("UNSLOTH_COMPILE_IGNORE_ERRORS", "1") == "1"
 # Just remove max_autotune_gemm warning
-import functools
 from torch._inductor.runtime.hints import DeviceProperties
 
 @functools.lru_cache(None)
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 77f953a5a3..15a506174f 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -591,20 +591,12 @@ class FastModel(FastBaseModel):
                 "if name.endswith(('q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj', 'head')): module.to(torch.float16); "\
                 "os.environ['TRITON_F32_DEFAULT'] = 'ieee';"
         elif "gpt-oss" in lowered_model_name:
-            os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1"
-	     # the temporary patches for init need UNSLOTH_MODEL_NAME to be set
-            # which doesn't happen at import so manually call here
-            # before creating the compiled cache
-            try:
-                from unsloth_zoo.temporary_patches.gpt_oss import (
-                    patch_GptOssExperts_MXFP4,
-                    patch_GptOssExperts_bitsandbytes,
-                )
-
-                patch_GptOssExperts_MXFP4()
-                patch_GptOssExperts_bitsandbytes()
-            except:
-                pass
+            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:
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 623f263127..b9cc01e530 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -365,8 +365,10 @@ class FastBaseModel:
             allow_float16_runs = (checker == "float16" and dtype == torch.float16)
 
             if allow_all_runs or allow_float16_runs:
-                dtype = eval(_dtype)
-                bnb_compute_dtype = eval(_bnb_compute_dtype)
+                if eval(_dtype) is not None:
+                    dtype = eval(_dtype)
+                if eval(_bnb_compute_dtype) is not None:
+                    bnb_compute_dtype = eval(_bnb_compute_dtype)
                 correct_dtype = bnb_compute_dtype
                 custom_datatype = _custom_datatype
                 # Execute code as well