Merge branch 'fix-issue-5344-quantization-guardrail' into feat-gemma4-moe-4bit-swap

unsloth/_gpu_init.py

@@ -209,7 +209,7 @@ del fix_peft_transformers_weight_conversion_import
 del patch_peft_weight_converter_compatibility
 
 # Torch 2.4 has including_emulation
-if DEVICE_TYPE == "cuda" and torch.cuda.is_available():
+if DEVICE_TYPE == "cuda":
     major_version, minor_version = torch.cuda.get_device_capability()
     SUPPORTS_BFLOAT16 = major_version >= 8
 
@@ -233,18 +233,12 @@ elif DEVICE_TYPE == "xpu":
     # torch.xpu.is_bf16_supported() does not have including_emulation
     # set SUPPORTS_BFLOAT16 as torch.xpu.is_bf16_supported()
     SUPPORTS_BFLOAT16 = torch.xpu.is_bf16_supported()
-else:
-    # CPU-only CI under UNSLOTH_ALLOW_CPU=1. We can't probe device
-    # capability, so assume no bf16 -- training won't run on this host
-    # anyway, this branch only exists to let `import unsloth.trainer`
-    # succeed for source-inspection tests.
-    SUPPORTS_BFLOAT16 = False
 
 # For Gradio HF Spaces?
 # if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ:
 import triton
 
-if DEVICE_TYPE == "cuda" and torch.cuda.is_available():
+if DEVICE_TYPE == "cuda":
     libcuda_dirs = lambda: None
     if Version(triton.__version__) >= Version("3.0.0"):
         try:

unsloth/device_type.py

@@ -52,6 +52,13 @@ def is_hip():
 
 @functools.cache
 def get_device_type():
+    # Test-only CPU fallback. Short-circuits the detection chain so the
+    # rest of the function -- and every DEVICE_TYPE == "cuda" branch in
+    # the codebase -- behaves identically to a real CUDA host. The env
+    # var is read exactly once per process because get_device_type is
+    # @functools.cache'd, so production hosts pay no runtime cost.
+    if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1":
+        return "cuda"
     if _IS_MLX:
         return "mlx"
     if hasattr(torch, "cuda") and torch.cuda.is_available():
@@ -63,10 +70,6 @@ def get_device_type():
     # Check torch.accelerator
     if hasattr(torch, "accelerator"):
         if not torch.accelerator.is_available():
-            # Test-only CPU fallback. The env var is read exactly once per
-            # process because get_device_type is @functools.cache'd.
-            if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1":
-                return "cuda"
             raise NotImplementedError(
                 "Unsloth cannot find any torch accelerator? You need a GPU."
             )
@@ -77,8 +80,6 @@ def get_device_type():
                 f"But `torch.accelerator.current_accelerator()` works with it being = `{accelerator}`\n"
                 f"Please reinstall torch - it's most likely broken :("
             )
-    if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1":
-        return "cuda"
     raise NotImplementedError(
         "Unsloth currently only works on NVIDIA, AMD and Intel GPUs."
     )

unsloth/models/_utils.py

@@ -1204,7 +1204,7 @@ SUPPORTS_BFLOAT16 = False
 HAS_FLASH_ATTENTION = False
 HAS_FLASH_ATTENTION_SOFTCAPPING = False
 
-if DEVICE_TYPE == "cuda" and torch.cuda.is_available():
+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)