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)