From 134638dd0ea321167c26ae5bd36b97e34b586f9b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 11:32:00 +0000 Subject: [PATCH] Harden ROCm detection, fix VRAM heuristic, and expand RDNA2 coverage - Windows ROCm detection: validate actual GPU presence via hipinfo/amd-smi output markers instead of just checking tool existence on PATH - _ensure_rocm_torch: validate nvidia-smi actually reports a GPU before giving NVIDIA precedence (fixes AMD-only hosts with stale NVIDIA tools) - amd.py _parse_numeric: handle dict-shaped metric objects from newer amd-smi versions ({"value": 10, "unit": "W"}) and strip MiB/GiB units - amd.py VRAM heuristic: raise threshold from 100k to 10M to correctly handle MI300X (192 GB = 196608 MB) and other high-VRAM GPUs - amd.py visible GPU: use AMD-reported GPU IDs instead of enumerate index so non-dense sets like CUDA_VISIBLE_DEVICES=1,3 report correctly - install.sh: add ROCm <6.0 minimum version guard (no PyTorch wheels exist for older versions); fix rocm7.1* glob to not match rocm7.10+ - is_rdna: add gfx1033-1036 for RDNA2 mobile GPUs (RX 6600M etc.) - worker.py: increase ROCm source build timeout from 600s to 1800s; fix success log message for ROCm source builds - Tests: update mocks for _has_usable_nvidia_gpu, add RDNA2 target asserts --- install.sh | 6 ++- studio/backend/core/training/worker.py | 9 ++-- studio/backend/utils/hardware/amd.py | 32 ++++++++++---- studio/install_llama_prebuilt.py | 23 ++++++---- studio/install_python_stack.py | 24 +++++++++-- tests/studio/install/test_rocm_support.py | 51 +++++++++++++++-------- unsloth/kernels/utils.py | 4 ++ 7 files changed, 108 insertions(+), 41 deletions(-) diff --git a/install.sh b/install.sh index 4d121bf331..9df0bb7347 100755 --- a/install.sh +++ b/install.sh @@ -1019,12 +1019,16 @@ get_torch_index_url() { *) _rocm_tag="" ;; # reject malformed (empty, garbled, or major=0) esac if [ -n "$_rocm_tag" ]; then + # Minimum supported: ROCm 6.0 (no PyTorch wheels exist for older) + case "$_rocm_tag" in + rocm[1-5].*) echo "$_base/cpu"; return ;; + esac # ROCm 7.2 only has torch 2.11.0 which exceeds current bounds (<2.11.0). # Fall back to rocm7.1 index which has torch 2.10.0. # TODO: uncomment the next line when torch upper bound is bumped to >=2.11.0 # echo "$_base/$_rocm_tag"; return case "$_rocm_tag" in - rocm6.*|rocm7.0*|rocm7.1*) + rocm6.*|rocm7.0|rocm7.0.*|rocm7.1|rocm7.1.*) echo "$_base/$_rocm_tag" ;; *) # ROCm 7.2+ (including future 10.x+): cap to rocm7.1 diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 8554173d01..d4a4b6eade 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -284,8 +284,8 @@ def _install_package_wheel_first( f"{pypi_name}=={pypi_version}", ] - # Source compilation on ROCm can take 5-10 minutes; use a generous timeout - timeout = 600 if is_hip else 300 + # Source compilation on ROCm can take 10-30 minutes; use a generous timeout + timeout = 1800 if is_hip else 300 try: result = _sp.run( @@ -331,7 +331,10 @@ def _install_package_wheel_first( ) return - logger.info("Installed %s from PyPI", display_name) + if is_hip: + logger.info("Compiled and installed %s from source for ROCm", display_name) + else: + logger.info("Installed %s from PyPI", display_name) def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None: diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index 8473ce8e51..f8e156e773 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -17,7 +17,7 @@ from loggers import get_logger logger = get_logger(__name__) -def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[dict]: +def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]: """Run amd-smi with the given arguments and return parsed JSON, or None.""" try: result = subprocess.run( @@ -43,11 +43,17 @@ def _parse_numeric(value: Any) -> Optional[float]: """Extract a numeric value from amd-smi output (may be str, int, float, or dict).""" if value is None: return None + # Newer amd-smi versions emit {"value": 10, "unit": "W"} + if isinstance(value, dict): + return _parse_numeric(value.get("value")) if isinstance(value, (int, float)): - return float(value) + import math + f = float(value) + return f if math.isfinite(f) else None if isinstance(value, str): - # Strip units like "W", "C", "%", "MB" etc. - cleaned = value.strip().rstrip("WCMBGb% ").strip() + # Strip units like "W", "C", "%", "MB", "MiB", "GB", "GiB" etc. + import re + cleaned = re.sub(r'\s*[A-Za-z/%]+$', '', value.strip()) if not cleaned or cleaned.lower() in ("n/a", "none", "unknown"): return None try: @@ -112,16 +118,18 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: vram_used_bytes = None vram_total_bytes = None - # Convert VRAM from bytes to MB if values are large (>10000 = likely bytes) + # Convert VRAM from bytes to MB if values are very large. + # amd-smi typically reports in MB, but some versions report bytes. + # Threshold: 10 million -- no GPU has <10 MB, and even 10 TB = 10M MB. vram_used_mb = None vram_total_mb = None if vram_used_bytes is not None: - if vram_used_bytes > 100000: # Likely bytes + if vram_used_bytes > 10_000_000: # Likely bytes (>10M) vram_used_mb = vram_used_bytes / (1024 * 1024) else: # Likely already MB vram_used_mb = vram_used_bytes if vram_total_bytes is not None: - if vram_total_bytes > 100000: # Likely bytes + if vram_total_bytes > 10_000_000: # Likely bytes (>10M) vram_total_mb = vram_total_bytes / (1024 * 1024) else: # Likely already MB vram_total_mb = vram_total_bytes @@ -211,12 +219,18 @@ def get_visible_gpu_utilization( "index_kind": "physical", } - gpu_list = data if isinstance(data, list) else data.get("gpus", [data]) + gpu_list = data if isinstance(data, list) else data.get("gpus", data.get("gpu", [data])) visible_set = set(parent_visible_ids) ordinal_map = {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)} devices = [] - for idx, gpu_data in enumerate(gpu_list): + for fallback_idx, gpu_data in enumerate(gpu_list): + # Use AMD-reported GPU ID when available, fall back to enumeration index + raw_id = gpu_data.get("gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx))) if isinstance(gpu_data, dict) else fallback_idx + try: + idx = int(raw_id) + except (TypeError, ValueError): + idx = fallback_idx if idx not in visible_set: continue metrics = _extract_gpu_metrics(gpu_data) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index abe31ebfff..2ea84eede5 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -1450,15 +1450,22 @@ def detect_host() -> HostInfo: has_rocm = True break elif is_windows: - # Windows: check for HIP runtime DLL or hipinfo tool - if shutil.which("hipinfo") or shutil.which("amd-smi"): - has_rocm = True - elif any( - Path(d).joinpath("amdhip64.dll").exists() - for d in os.environ.get("PATH", "").split(os.pathsep) - if d + # Windows: validate actual AMD GPU presence (not just tool/DLL existence) + for _cmd, _marker in ( + (["hipinfo"], "gcnarchname"), + (["amd-smi", "list"], "gpu"), ): - has_rocm = True + _exe = shutil.which(_cmd[0]) + if not _exe: + continue + try: + _result = run_capture([_exe, *_cmd[1:]], timeout = 10) + except Exception: + continue + if _result.returncode == 0 and _result.stdout.strip(): + if _marker in _result.stdout.lower(): + has_rocm = True + break return HostInfo( system = system, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index dc6b3fe49e..cba37886a2 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -109,6 +109,24 @@ def _has_rocm_gpu() -> bool: return False +def _has_usable_nvidia_gpu() -> bool: + """Return True only when nvidia-smi exists AND reports at least one GPU.""" + exe = shutil.which("nvidia-smi") + if not exe: + return False + try: + result = subprocess.run( + [exe, "-L"], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + timeout=10, + ) + except Exception: + return False + return result.returncode == 0 and "GPU " in result.stdout + + def _ensure_rocm_torch() -> None: """Reinstall torch with ROCm wheels when the venv received CPU-only torch. @@ -118,8 +136,8 @@ def _ensure_rocm_torch() -> None: (NVIDIA takes precedence). Uses pip_install() to respect uv, constraints, and --python targeting. """ - # NVIDIA takes precedence on mixed hosts - if shutil.which("nvidia-smi"): + # NVIDIA takes precedence on mixed hosts -- but only if an actual GPU is usable + if _has_usable_nvidia_gpu(): return rocm_root = os.environ.get("ROCM_PATH") or "/opt/rocm" if not os.path.isdir(rocm_root) and not shutil.which("hipcc"): @@ -707,7 +725,7 @@ def install_python_stack() -> int: # Windows + AMD GPU: PyTorch does not publish ROCm wheels for Windows. # Detect and warn so users know manual steps are needed for GPU training. - if IS_WINDOWS and not NO_TORCH: + if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu(): if shutil.which("hipinfo") or shutil.which("amd-smi"): _safe_print( _dim(" Note:"), diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index ff71da6b25..eaa1ce2986 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -49,6 +49,7 @@ _STACK_SPEC.loader.exec_module(stack_mod) _detect_rocm_version = stack_mod._detect_rocm_version _ensure_rocm_torch = stack_mod._ensure_rocm_torch _has_rocm_gpu = stack_mod._has_rocm_gpu +_has_usable_nvidia_gpu = stack_mod._has_usable_nvidia_gpu _ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX @@ -396,20 +397,20 @@ class TestHostInfoRocm: assert host.has_rocm is False assert host.has_usable_nvidia is True - def test_detect_host_with_rocm_path_env(self): - """detect_host() checks ROCM_PATH env var for ROCm detection.""" - # Verify the detect_host function source references ROCM_PATH + def test_detect_host_has_rocm_detection_logic(self): + """detect_host() should have ROCm GPU detection logic.""" import inspect source = inspect.getsource(prebuilt_mod.detect_host) - assert "ROCM_PATH" in source or "rocm" in source.lower() + # Must probe for actual GPU, not just tool presence + assert "rocminfo" in source or "amd-smi" in source def test_detect_host_windows_rocm_detection(self): - """detect_host() source should have Windows-specific HIP detection.""" + """detect_host() source should have Windows-specific ROCm GPU detection.""" import inspect source = inspect.getsource(prebuilt_mod.detect_host) - assert "hipinfo" in source or "amdhip64" in source + assert "hipinfo" in source or "amd-smi" in source # ============================================================================= @@ -520,7 +521,8 @@ class TestEnsureRocmTorch: """Verify ROCm torch reinstall logic.""" @patch.object(stack_mod, "pip_install") - def test_no_rocm_skips(self, mock_pip): + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + def test_no_rocm_skips(self, mock_nvidia, mock_pip): """No ROCm toolchain should skip entirely.""" with patch("os.path.isdir", return_value = False): with patch("shutil.which", return_value = None): @@ -528,9 +530,10 @@ class TestEnsureRocmTorch: mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) - def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_pip): + def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """If torch already has CUDA, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -541,9 +544,10 @@ class TestEnsureRocmTorch: mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) - def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_pip): + def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """If torch already has HIP, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -554,9 +558,10 @@ class TestEnsureRocmTorch: mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) - def test_cpu_torch_gets_rocm_reinstall(self, mock_ver, mock_gpu, mock_pip): + def test_cpu_torch_gets_rocm_reinstall(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """CPU-only torch on ROCm host should trigger reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -572,9 +577,10 @@ class TestEnsureRocmTorch: assert "bitsandbytes" in str(bnb_call) @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 3)) - def test_rocm_63_selects_correct_tag(self, mock_ver, mock_gpu, mock_pip): + def test_rocm_63_selects_correct_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """ROCm 6.3 should select rocm6.3 tag.""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -586,9 +592,10 @@ class TestEnsureRocmTorch: assert "rocm6.3" in str(torch_call) @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (5, 0)) - def test_old_rocm_skips(self, mock_ver, mock_gpu, mock_pip): + def test_old_rocm_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """ROCm version too old (below 6.0) should skip.""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -599,10 +606,11 @@ class TestEnsureRocmTorch: mock_pip.assert_not_called() @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = None) def test_version_unreadable_prints_warning( - self, mock_ver, mock_gpu, mock_pip, capsys + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, capsys ): """ROCm detected but version unreadable should print warning and skip.""" with patch("os.path.isdir", return_value = True): @@ -612,9 +620,10 @@ class TestEnsureRocmTorch: assert "unreadable" in captured.out @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) - def test_rocm_72_selects_71_tag(self, mock_ver, mock_gpu, mock_pip): + def test_rocm_72_selects_71_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """ROCm 7.2 should select rocm7.1 tag (capped, not in mapping).""" mock_probe = MagicMock() mock_probe.returncode = 0 @@ -626,9 +635,10 @@ class TestEnsureRocmTorch: assert "rocm7.1" in str(torch_call) @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) - def test_probe_timeout_triggers_reinstall(self, mock_ver, mock_gpu, mock_pip): + def test_probe_timeout_triggers_reinstall(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): """Probe subprocess timeout should not crash; should proceed to reinstall.""" with patch("os.path.isdir", return_value = True): with patch( @@ -640,8 +650,9 @@ class TestEnsureRocmTorch: assert "rocm7.1" in str(mock_pip.call_args_list[0]) @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = False) - def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_pip): + def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip): """ROCm tools present but no actual AMD GPU should skip entirely.""" with patch("os.path.isdir", return_value = True): _ensure_rocm_torch() @@ -865,7 +876,9 @@ class TestInstallShStructure: source = sh_path.read_text() assert 'echo "$_base/rocm7.1"' in source # fallback for unknown versions # Allowlisted versions should pass through directly - assert "rocm6.*|rocm7.0*|rocm7.1*)" in source + assert "rocm6.*" in source + assert "rocm7.0" in source + assert "rocm7.1" in source def test_rocm_tag_validation_guard_exists(self): """install.sh should validate _rocm_tag with a case guard.""" @@ -1269,6 +1282,10 @@ class TestIsRdnaExpansion: assert "gfx1030" in func_body assert "gfx1031" in func_body assert "gfx1032" in func_body + assert "gfx1033" in func_body + assert "gfx1034" in func_body + assert "gfx1035" in func_body + assert "gfx1036" in func_body def test_is_rdna_source_has_rdna3(self): """is_rdna() should include RDNA3 architectures.""" diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 049afac96d..09b03a597b 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -94,6 +94,10 @@ def is_rdna(): "gfx1030", "gfx1031", "gfx1032", + "gfx1033", + "gfx1034", + "gfx1035", + "gfx1036", # RDNA3 (Navi 31-33) "gfx1100", "gfx1101",