diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index aaa0452406..b9b61cdcfe 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -9,6 +9,7 @@ from .hardware import ( DeviceType, DEVICE, CHAT_ONLY, + IS_ROCM, detect_hardware, get_device, is_apple_silicon, @@ -49,6 +50,7 @@ __all__ = [ "DeviceType", "DEVICE", "CHAT_ONLY", + "IS_ROCM", "detect_hardware", "get_device", "is_apple_silicon", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 742e8f6b7e..244ef3fe98 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -43,6 +43,7 @@ class DeviceType(str, Enum): DEVICE: Optional[DeviceType] = None CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.) +IS_ROCM: bool = False # True when running on AMD ROCm (HIP) -- display/logging only # ========== Detection ========== @@ -85,10 +86,11 @@ def detect_hardware() -> DeviceType: 2. MLX (Apple Silicon via MLX framework) 3. CPU (fallback) """ - global DEVICE, CHAT_ONLY - CHAT_ONLY = True # reset -- only CUDA sets it to False + global DEVICE, CHAT_ONLY, IS_ROCM + CHAT_ONLY = True # reset -- only CUDA/ROCm sets it to False + IS_ROCM = False - # --- CUDA: try PyTorch --- + # --- CUDA / ROCm: try PyTorch --- if _has_torch(): import torch @@ -96,7 +98,16 @@ def detect_hardware() -> DeviceType: DEVICE = DeviceType.CUDA CHAT_ONLY = False device_name = torch.cuda.get_device_properties(0).name - print(f"Hardware detected: CUDA — {device_name}") + + # Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes. + # DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP. + if getattr(torch.version, "hip", None) is not None: + IS_ROCM = True + print( + f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}" + ) + else: + print(f"Hardware detected: CUDA -- {device_name}") return DEVICE # --- XPU: Intel GPU --- @@ -315,13 +326,15 @@ def get_package_versions() -> Dict[str, Optional[str]]: except PackageNotFoundError: versions[name] = None - # CUDA toolkit version bundled with torch + # GPU runtime version bundled with torch try: import torch versions["cuda"] = getattr(torch.version, "cuda", None) + versions["rocm"] = getattr(torch.version, "hip", None) except Exception: versions["cuda"] = None + versions["rocm"] = None return versions diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 8be6bb5a5a..07949cd32e 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -1103,7 +1103,7 @@ def patch_sft_trainer_tokenizer(): " a = np.array([int(x.decode('utf-8'))/1024 for x in a])\n" "except:\n" " if not torch.cuda.is_available():\n" - " raise RuntimeError('Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!')\n" + " raise RuntimeError('Unsloth: No GPU detected. AMD ROCm users: install ROCm-enabled PyTorch -- see https://docs.unsloth.ai/get-started/install-and-update/amd')\n" "if ((a - PRE_CHECK) >= 1).sum() > 1:\n" " raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!')\n" "for _ in range(3):\n"