From 8028c526002ae2991a7383efbf6c97414381cd70 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 31 Mar 2026 08:31:31 +0000
Subject: [PATCH] Add IS_ROCM flag to hardware module and fix AMD error message

Hardware module changes:
- Add IS_ROCM global flag (bool, default False) to hardware.py
- detect_hardware() sets IS_ROCM=True when torch.version.hip is set,
  while keeping DeviceType as CUDA (torch.cuda.* works on ROCm via HIP)
- Print "ROCm (HIP X.Y)" instead of just "CUDA" when IS_ROCM is True
- get_package_versions() now returns a "rocm" key with the HIP version
  (torch.version.hip) alongside the existing "cuda" key
- Export IS_ROCM from studio/backend/utils/hardware/__init__.py

Tokenizer error message fix:
- Replace "We do not support AMD" with a helpful message pointing to
  ROCm installation docs at docs.unsloth.ai, since AMD is now supported
---
 studio/backend/utils/hardware/__init__.py |  2 ++
 studio/backend/utils/hardware/hardware.py | 21 ++++++++++++++++-----
 unsloth/tokenizer_utils.py                |  2 +-
 3 files changed, 19 insertions(+), 6 deletions(-)

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..d5de59a592 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,14 @@ 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 +324,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"