From 77e09297353c36dac34af6f1f5688611d80902d0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 15 May 2026 19:41:09 -0700 Subject: [PATCH] revert: stop touching DEVICE_TYPE == "cuda" branches for CPU CI (#5473) #5429 (cb15a7a5) tightened three production-path branches to DEVICE_TYPE == "cuda" and torch.cuda.is_available() and added a new else: SUPPORTS_BFLOAT16 = False arm to let `import unsloth.trainer` survive on a CPU-only CI host. We already ship the package on Intel XPU / AMD HIP / NVIDIA CUDA and don't want any extra branching in those hot paths. Move the entire CPU-CI handling to one place -- the top of unsloth/device_type.get_device_type() -- so the UNSLOTH_ALLOW_CPU=1 sentinel short-circuits detection and returns "cuda" before any torch probe runs. Every downstream DEVICE_TYPE == "cuda" branch then behaves identically to a real CUDA host, with no additional checks. The two existing duplicate UNSLOTH_ALLOW_CPU returns later in the function are dropped (the new top-of-function check covers both). Revert the three call-site changes: - unsloth/_gpu_init.py:212 -> back to `if DEVICE_TYPE == "cuda":` - unsloth/_gpu_init.py:247 -> back to `if DEVICE_TYPE == "cuda":` - unsloth/models/_utils.py:1207 -> back to `if DEVICE_TYPE == "cuda":` - unsloth/_gpu_init.py: drop the new `else: SUPPORTS_BFLOAT16 = False` branch (dead under the top-of-function short-circuit). Keep the two env-var gates that are needed for zoo's drift detectors to inspect pristine TRL source (no behavioural change on production hosts that never set UNSLOTH_ALLOW_CPU): - unsloth/_gpu_init.py: `if env != "1": _patch_trl_trainer()` - unsloth/models/rl.py:PatchFastRL: `if env == "1": return` Verified: - CUDA_VISIBLE_DEVICES=5 python -c "import unsloth.trainer" produces UnslothSFTTrainer.__init__ (TRL still patched on real hosts). - UNSLOTH_ALLOW_CPU=1 + aggressive cuda spoof import succeeds and trl.SFTTrainer.__init__.__qualname__ stays SFTTrainer.__init__. - pytest tests/_zoo_compiler_cache_shim.py -> 5 passed, 1 skipped. --- unsloth/_gpu_init.py | 10 ++-------- unsloth/device_type.py | 13 +++++++------ unsloth/models/_utils.py | 2 +- 3 files changed, 10 insertions(+), 15 deletions(-) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index a30111b529..df446195fb 100644 --- a/unsloth/_gpu_init.py +++ b/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: diff --git a/unsloth/device_type.py b/unsloth/device_type.py index 6a82e42e8c..f7a330594b 100644 --- a/unsloth/device_type.py +++ b/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." ) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 7bf8866f46..5c3a5742e4 100644 --- a/unsloth/models/_utils.py +++ b/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)