From 0539621fa15969d9f3922665f8a7115d2d6856d2 Mon Sep 17 00:00:00 2001
From: LeoBorcherding <borchborchmail@gmail.com>
Date: Tue, 5 May 2026 13:08:38 -0500
Subject: [PATCH] fix(studio): set HIP_VISIBLE_DEVICES in apply_gpu_ids for
 ROCm training workers

Training workers are spawned via multiprocessing spawn before detect_hardware()
runs, so IS_ROCM is still False. If the user never set HIP_VISIBLE_DEVICES in
their shell, _inherits_rocm_visibility is also False, leaving the worker with
only CUDA_VISIBLE_DEVICES set. On ROCm hosts the HIP runtime honors
HIP_VISIBLE_DEVICES over CUDA_VISIBLE_DEVICES, so the worker saw the full
device list and torch raised "no usable HIP accelerator" on some setups.

Fall back to probing torch.version.hip (a build-time attribute, safe to read
before GPU init) to detect ROCm when neither IS_ROCM nor inherited env vars
are available. Mirrors the existing fix in llama_cpp.py for llama-server
subprocess GPU pinning.

Fixes https://github.com/unslothai/unsloth/issues/5180
---
 studio/backend/utils/hardware/hardware.py | 15 +++++++++--
 tests/studio/install/test_rocm_support.py | 33 +++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index be31c00a78..7246a0519a 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -1391,14 +1391,25 @@ def apply_gpu_ids(gpu_ids) -> None:
     # parent process already set a ROCm visibility variable -- that
     # way a downstream ROCm process inherits the narrowed mask even
     # before Studio's hardware detection has classified the host.
+    # As a final fallback, probe torch.version.hip directly so spawned
+    # training workers on AMD hosts where the user never set HIP_VISIBLE_DEVICES
+    # still get the correct ROCm visibility mask (mirrors the llama_cpp.py
+    # approach for llama-server subprocess GPU pinning).
     _inherits_rocm_visibility = (
         "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
     )
-    if IS_ROCM or _inherits_rocm_visibility:
+    _is_rocm = IS_ROCM or _inherits_rocm_visibility
+    if not _is_rocm:
+        try:
+            import torch as _torch
+            _is_rocm = bool(getattr(_torch.version, "hip", None))
+        except Exception:
+            pass
+    if _is_rocm:
         os.environ["HIP_VISIBLE_DEVICES"] = value
         os.environ["ROCR_VISIBLE_DEVICES"] = value
     _visible_gpu_count = None
-    if IS_ROCM or _inherits_rocm_visibility:
+    if _is_rocm:
         logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value)
     else:
         logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value)
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index 48831fd57b..9ce6e53998 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -1250,6 +1250,39 @@ class TestHardwareAmdBranching:
         assert "amd.get_physical_gpu_count" in func_body
 
 
+# =============================================================================
+# TEST: hardware.py -- apply_gpu_ids ROCm fallback (issue #5180)
+# =============================================================================
+
+
+class TestApplyGpuIdsRocmFallback:
+    """Verify apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when
+    IS_ROCM is still False (worker subprocess before detect_hardware runs)."""
+
+    def test_apply_gpu_ids_source_checks_torch_version_hip(self):
+        """apply_gpu_ids should fall back to torch.version.hip when IS_ROCM is False."""
+        hw_path = (
+            PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
+        )
+        source = hw_path.read_text()
+        func_start = source.find("def apply_gpu_ids")
+        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
+        assert 'getattr(_torch.version, "hip", None)' in func_body or \
+               "getattr(torch.version, 'hip', None)" in func_body or \
+               "torch.version.hip" in func_body
+
+    def test_apply_gpu_ids_source_sets_hip_visible_devices(self):
+        """apply_gpu_ids should set HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES."""
+        hw_path = (
+            PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
+        )
+        source = hw_path.read_text()
+        func_start = source.find("def apply_gpu_ids")
+        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
+        assert "HIP_VISIBLE_DEVICES" in func_body
+        assert "ROCR_VISIBLE_DEVICES" in func_body
+
+
 # =============================================================================
 # TEST: install_python_stack.py -- Windows AMD warning
 # =============================================================================