GPT OSS fixes

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)

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:

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