From f17e007cafac4c39b88dfad6769a413798db5dac Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 11:08:20 +0000 Subject: [PATCH] Add AMD ROCm gaps: Mamba/SSM source builds, GPU monitoring, Windows messaging, RDNA expansion - worker.py: Add HIP detection to causal-conv1d/mamba-ssm probe, check for hipcc before ROCm source builds, improve status messages and error reporting, add timeout and uv support for the source build fallback - amd.py: New AMD GPU monitoring module via amd-smi metric --json, mirroring nvidia.py structure (utilization, temperature, power, VRAM) - hardware.py: Branch to amd.py when IS_ROCM is True for GPU utilization, visible GPU queries, and physical GPU count - install_python_stack.py: Detect AMD GPUs on Windows and warn that ROCm-enabled PyTorch must be installed manually - kernels/utils.py: Expand is_rdna() to cover RDNA2 (gfx1030-1032), RDNA3 (gfx1102-1103), RDNA3.5 (gfx1150-1152) alongside existing entries - tests: Add 32 new tests covering all changes (95/95 pass) --- studio/backend/core/training/worker.py | 100 ++++-- studio/backend/utils/hardware/amd.py | 224 +++++++++++++ studio/backend/utils/hardware/hardware.py | 96 ++++-- studio/install_python_stack.py | 13 + tests/studio/install/test_rocm_support.py | 367 ++++++++++++++++++++++ unsloth/kernels/utils.py | 14 +- 6 files changed, 765 insertions(+), 49 deletions(-) create mode 100644 studio/backend/utils/hardware/amd.py diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 699cfe74f7..187f43b3e0 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -86,6 +86,7 @@ def _probe_causal_conv1d_env() -> dict[str, str] | None: "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', " "'torch_mm': torch_mm, " "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', " + "'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', " "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()" "}))" ), @@ -237,25 +238,88 @@ def _install_package_wheel_first( else: logger.info("No published %s wheel found: %s", display_name, wheel_url) - _send_status(event_queue, f"Installing {display_name} from PyPI...") - pypi_cmd = [ - sys.executable, - "-m", - "pip", - "install", - "--no-build-isolation", - "--no-deps", - "--no-cache-dir", - f"{pypi_name}=={pypi_version}", - ] - result = _sp.run( - pypi_cmd, - stdout = _sp.PIPE, - stderr = _sp.STDOUT, - text = True, - ) + is_hip = env and env.get("hip_version") + if is_hip and not shutil.which("hipcc"): + logger.error( + "%s requires hipcc for source compilation on ROCm. " + "Install the ROCm HIP SDK: https://rocm.docs.amd.com", + display_name, + ) + _send_status( + event_queue, + f"{display_name}: hipcc not found (ROCm HIP SDK required)", + ) + return + + if is_hip: + _send_status( + event_queue, + f"Compiling {display_name} from source for ROCm " + "(this may take several minutes)...", + ) + else: + _send_status(event_queue, f"Installing {display_name} from PyPI...") + + # Prefer uv for faster dependency resolution when available + if shutil.which("uv"): + pypi_cmd = [ + "uv", "pip", "install", + "--python", sys.executable, + "--no-build-isolation", + "--no-deps", + f"{pypi_name}=={pypi_version}", + ] + else: + pypi_cmd = [ + sys.executable, + "-m", + "pip", + "install", + "--no-build-isolation", + "--no-deps", + "--no-cache-dir", + 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 + + try: + result = _sp.run( + pypi_cmd, + stdout = _sp.PIPE, + stderr = _sp.STDOUT, + text = True, + timeout = timeout, + ) + except _sp.TimeoutExpired: + logger.error( + "%s installation timed out after %ds", display_name, timeout, + ) + _send_status( + event_queue, + f"{display_name} installation timed out after {timeout}s", + ) + return + if result.returncode != 0: - logger.error("Failed to install %s from PyPI:\n%s", display_name, result.stdout) + if is_hip: + # Surface a clear error for ROCm source build failures + error_lines = (result.stdout or "").strip().splitlines() + snippet = "\n".join(error_lines[-5:]) if error_lines else "(no output)" + logger.error( + "Failed to compile %s for ROCm:\n%s", display_name, result.stdout, + ) + _send_status( + event_queue, + f"Failed to compile {display_name} for ROCm. " + "Check that hipcc and ROCm development headers are installed.\n" + f"{snippet}", + ) + else: + logger.error( + "Failed to install %s from PyPI:\n%s", display_name, result.stdout, + ) return logger.info("Installed %s from PyPI", display_name) diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py new file mode 100644 index 0000000000..df4f83d5ca --- /dev/null +++ b/studio/backend/utils/hardware/amd.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""AMD GPU monitoring via amd-smi. + +Mirrors the nvidia.py module structure so hardware.py can swap backends +based on IS_ROCM. All functions return the same dict shapes as their +nvidia.py counterparts. +""" + +import json +import subprocess +from typing import Any, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + + +def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[dict]: + """Run amd-smi with the given arguments and return parsed JSON, or None.""" + try: + result = subprocess.run( + ["amd-smi", *args, "--json"], + capture_output=True, + text=True, + timeout=timeout, + ) + except (OSError, subprocess.TimeoutExpired) as e: + logger.warning("amd-smi query failed: %s", e) + return None + if result.returncode != 0 or not result.stdout.strip(): + logger.warning("amd-smi returned code %d", result.returncode) + return None + try: + return json.loads(result.stdout) + except json.JSONDecodeError: + logger.warning("Failed to parse amd-smi JSON output") + return None + + +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 + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + # Strip units like "W", "C", "%", "MB" etc. + cleaned = value.strip().rstrip("WCMBGb% ").strip() + if not cleaned or cleaned.lower() in ("n/a", "none", "unknown"): + return None + try: + return float(cleaned) + except (ValueError, TypeError): + return None + return None + + +def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: + """Extract standardized metrics from a single GPU's amd-smi data.""" + # amd-smi metric output structure varies by version; try common paths + usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {})) + if isinstance(usage, dict): + gpu_util = _parse_numeric(usage.get("gfx_activity", usage.get("gpu_use_percent"))) + else: + gpu_util = _parse_numeric(usage) + + # Temperature + temp_data = gpu_data.get("temperature", {}) + if isinstance(temp_data, dict): + temp = _parse_numeric( + temp_data.get("edge", temp_data.get("temperature_edge", + temp_data.get("hotspot", temp_data.get("temperature_hotspot")))) + ) + else: + temp = _parse_numeric(temp_data) + + # Power + power_data = gpu_data.get("power", {}) + if isinstance(power_data, dict): + power_draw = _parse_numeric( + power_data.get("current_socket_power", + power_data.get("average_socket_power", + power_data.get("socket_power"))) + ) + power_limit = _parse_numeric( + power_data.get("power_cap", power_data.get("max_power_limit")) + ) + else: + power_draw = None + power_limit = None + + # VRAM + vram_data = gpu_data.get("vram", gpu_data.get("fb_memory_usage", {})) + if isinstance(vram_data, dict): + vram_used_bytes = _parse_numeric( + vram_data.get("vram_used", vram_data.get("used")) + ) + vram_total_bytes = _parse_numeric( + vram_data.get("vram_total", vram_data.get("total")) + ) + else: + vram_used_bytes = None + vram_total_bytes = None + + # Convert VRAM from bytes to MB if values are large (>10000 = likely bytes) + vram_used_mb = None + vram_total_mb = None + if vram_used_bytes is not None: + if vram_used_bytes > 100000: # Likely bytes + 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 + vram_total_mb = vram_total_bytes / (1024 * 1024) + else: # Likely already MB + vram_total_mb = vram_total_bytes + + # Build the standardized dict (same shape as nvidia._build_gpu_metrics) + vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None + vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None + vram_util = ( + round((vram_used_mb / vram_total_mb) * 100, 1) + if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0 + else None + ) + power_util = ( + round((power_draw / power_limit) * 100, 1) + if power_draw is not None and power_limit and power_limit > 0 + else None + ) + + return { + "gpu_utilization_pct": gpu_util, + "temperature_c": temp, + "vram_used_gb": vram_used_gb, + "vram_total_gb": vram_total_gb, + "vram_utilization_pct": vram_util, + "power_draw_w": power_draw, + "power_limit_w": power_limit, + "power_utilization_pct": power_util, + } + + +def get_physical_gpu_count() -> Optional[int]: + """Return physical AMD GPU count via amd-smi, or None on failure.""" + data = _run_amd_smi("list") + if data is None: + return None + if isinstance(data, list): + return len(data) + # Some versions return a dict with a "gpu" key + gpus = data.get("gpu", data.get("gpus", [])) + if isinstance(gpus, list): + return len(gpus) + return None + + +def get_primary_gpu_utilization() -> dict[str, Any]: + """Return utilization metrics for the primary AMD GPU.""" + data = _run_amd_smi("metric", "-g", "0") + if data is None: + return {"available": False} + + # amd-smi may return a list with one entry or a dict + if isinstance(data, list): + if len(data) == 0: + return {"available": False} + gpu_data = data[0] + else: + gpu_data = data + + metrics = _extract_gpu_metrics(gpu_data) + metrics["available"] = True + return metrics + + +def get_visible_gpu_utilization( + parent_visible_ids: Optional[list[int]], + parent_cuda_visible_devices: Optional[str] = None, +) -> dict[str, Any]: + """Return utilization metrics for visible AMD GPUs.""" + if parent_visible_ids is None: + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "unresolved", + } + + data = _run_amd_smi("metric") + if data is None: + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + + gpu_list = data if isinstance(data, list) else data.get("gpus", [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): + if idx not in visible_set: + continue + metrics = _extract_gpu_metrics(gpu_data) + metrics["index"] = idx + metrics["index_kind"] = "physical" + metrics["visible_ordinal"] = ordinal_map.get(idx, len(devices)) + devices.append(metrics) + + return { + "available": len(devices) > 0, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": devices, + "index_kind": "physical", + } diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 7a3fad5994..049e250237 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -405,15 +405,26 @@ def get_gpu_utilization() -> Dict[str, Any]: device = get_device() if device == DeviceType.CUDA: - try: - from . import nvidia + if IS_ROCM: + try: + from . import amd - result = nvidia.get_primary_gpu_utilization() - if result.get("available"): - result["backend"] = device.value - return result - except Exception as e: - logger.warning("nvidia-smi utilization query failed: %s", e) + result = amd.get_primary_gpu_utilization() + if result.get("available"): + result["backend"] = device.value + return result + except Exception as e: + logger.warning("amd-smi utilization query failed: %s", e) + else: + try: + from . import nvidia + + result = nvidia.get_primary_gpu_utilization() + if result.get("available"): + result["backend"] = device.value + return result + except Exception as e: + logger.warning("nvidia-smi utilization query failed: %s", e) mem = get_gpu_memory_info() if device != DeviceType.CPU and mem.get("available"): @@ -438,18 +449,32 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: if device == DeviceType.CUDA: parent_visible_spec = _get_parent_visible_gpu_spec() - try: - from . import nvidia + if IS_ROCM: + try: + from . import amd - result = nvidia.get_visible_gpu_utilization( - parent_visible_spec["numeric_ids"], - parent_cuda_visible_devices = parent_visible_spec["raw"], - ) - if result.get("available"): - result["backend"] = device.value - return result - except Exception as e: - logger.warning("nvidia-smi visible GPU utilization query failed: %s", e) + result = amd.get_visible_gpu_utilization( + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result.get("available"): + result["backend"] = device.value + return result + except Exception as e: + logger.warning("amd-smi visible GPU utilization query failed: %s", e) + else: + try: + from . import nvidia + + result = nvidia.get_visible_gpu_utilization( + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result.get("available"): + result["backend"] = device.value + return result + except Exception as e: + logger.warning("nvidia-smi visible GPU utilization query failed: %s", e) # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) if device in (DeviceType.CUDA, DeviceType.XPU): @@ -1121,16 +1146,27 @@ def get_physical_gpu_count() -> int: device = get_device() if device == DeviceType.CUDA: - try: - from . import nvidia + if IS_ROCM: + try: + from . import amd - count = nvidia.get_physical_gpu_count() - if count is not None: - _physical_gpu_count = count - return _physical_gpu_count - except Exception: - pass - # nvidia-smi unavailable or failed — fall back to torch + count = amd.get_physical_gpu_count() + if count is not None: + _physical_gpu_count = count + return _physical_gpu_count + except Exception: + pass + else: + try: + from . import nvidia + + count = nvidia.get_physical_gpu_count() + if count is not None: + _physical_gpu_count = count + return _physical_gpu_count + except Exception: + pass + # SMI tool unavailable or failed -- fall back to torch count = _torch_get_physical_gpu_count() _physical_gpu_count = count if count is not None else 1 return _physical_gpu_count @@ -1153,8 +1189,8 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: device = get_device() if device in (DeviceType.CUDA, DeviceType.XPU): parent_visible_ids = get_parent_visible_gpu_ids() - # Try nvidia-smi first (NVIDIA only) - if device == DeviceType.CUDA: + # Try native SMI tool first (nvidia-smi for NVIDIA, skipped for ROCm) + if device == DeviceType.CUDA and not IS_ROCM: try: from . import nvidia diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 461f83d9b9..dc6b3fe49e 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -705,6 +705,19 @@ def install_python_stack() -> int: _progress("ROCm torch check") _ensure_rocm_torch() + # 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 shutil.which("hipinfo") or shutil.which("amd-smi"): + _safe_print( + _dim(" Note:"), + "AMD GPU detected on Windows. ROCm-enabled PyTorch must be", + ) + _safe_print( + " " * 8, + "installed manually. See: https://docs.unsloth.ai/get-started/install-and-update/amd", + ) + # 3. Extra dependencies _progress("unsloth extras") pip_install( diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index fe76d920a5..59c51da79a 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -6,6 +6,7 @@ All tests use mocks -- no AMD hardware required. """ import importlib.util +import json import os import subprocess import sys @@ -958,5 +959,371 @@ class TestLiveRegression: assert "cu1" in url or "cuda" in url.lower(), f"Expected CUDA URL, got: {url}" +# ============================================================================= +# TEST: worker.py -- ROCm Mamba/SSM source build path +# ============================================================================= + +# Load worker.py module +_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" + + +class TestWorkerRocmMambaSsm: + """Verify worker.py Mamba/SSM install logic on ROCm.""" + + def test_probe_returns_hip_version_field(self): + """_probe_causal_conv1d_env probe script should include hip_version.""" + source = _WORKER_PATH.read_text() + assert "hip_version" in source + + def test_probe_script_has_getattr_hip(self): + """Probe script should use getattr for torch.version.hip (safe on CUDA).""" + source = _WORKER_PATH.read_text() + assert "getattr(torch.version, 'hip', None)" in source + + def test_direct_wheel_url_returns_none_without_cuda_major(self): + """_direct_wheel_url should return None when cuda_major is empty (ROCm).""" + # Load module for function access + _worker_spec = importlib.util.spec_from_file_location( + "test_worker", _WORKER_PATH + ) + assert _worker_spec is not None and _worker_spec.loader is not None + worker_mod = importlib.util.module_from_spec(_worker_spec) + + # Mock all the imports worker.py needs + sys.modules["structlog"] = MagicMock() + sys.modules["loggers"] = MagicMock() + sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock()) + sys.modules["utils"] = MagicMock() + sys.modules["utils.hardware"] = MagicMock() + + try: + _worker_spec.loader.exec_module(worker_mod) + except Exception: + pytest.skip("Could not load worker module in test environment") + + env_rocm = { + "python_tag": "cp312", + "torch_mm": "2.6", + "cuda_major": "", + "hip_version": "7.1.12345", + "cxx11abi": "TRUE", + } + result = worker_mod._direct_wheel_url( + filename_prefix="causal_conv1d", + package_version="1.6.1", + release_tag="v1.6.1.post4", + release_base_url="https://github.com/Dao-AILab/causal-conv1d/releases/download", + env=env_rocm, + ) + assert result is None + + def test_hipcc_check_exists_in_source(self): + """worker.py should check for hipcc before ROCm source builds.""" + source = _WORKER_PATH.read_text() + assert "hipcc" in source + + def test_rocm_source_build_status_message(self): + """worker.py should send a specific status for ROCm source compilation.""" + source = _WORKER_PATH.read_text() + assert "Compiling" in source and "from source for ROCm" in source + + def test_rocm_build_failure_message(self): + """worker.py should send a clear error on ROCm build failure.""" + source = _WORKER_PATH.read_text() + assert "Failed to compile" in source and "for ROCm" in source + + def test_timeout_on_install(self): + """worker.py should have a timeout on pip install subprocess.""" + source = _WORKER_PATH.read_text() + assert "TimeoutExpired" in source + assert "timeout" in source + + +# ============================================================================= +# TEST: amd.py -- AMD GPU monitoring +# ============================================================================= + + +class TestAmdGpuMonitoring: + """Verify amd.py module structure and mock behavior.""" + + def test_amd_py_exists(self): + """amd.py should exist in the hardware directory.""" + amd_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" + ) + assert amd_path.exists() + + def test_amd_py_has_required_functions(self): + """amd.py should export the same function signatures as nvidia.py.""" + amd_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" + ) + source = amd_path.read_text() + assert "def get_physical_gpu_count" in source + assert "def get_primary_gpu_utilization" in source + assert "def get_visible_gpu_utilization" in source + + def test_amd_smi_json_parsing(self): + """Verify _extract_gpu_metrics parses amd-smi JSON correctly.""" + amd_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" + ) + _amd_spec = importlib.util.spec_from_file_location("test_amd", amd_path) + assert _amd_spec is not None and _amd_spec.loader is not None + amd_mod = importlib.util.module_from_spec(_amd_spec) + + sys.modules["loggers"] = MagicMock() + sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock()) + + try: + _amd_spec.loader.exec_module(amd_mod) + except Exception: + pytest.skip("Could not load amd module in test environment") + + # Simulate amd-smi metric JSON output + gpu_data = { + "usage": {"gfx_activity": "85"}, + "temperature": {"edge": "72"}, + "power": { + "current_socket_power": "200.5", + "power_cap": "300", + }, + "vram": { + "vram_used": 8192, # MB + "vram_total": 16384, # MB + }, + } + metrics = amd_mod._extract_gpu_metrics(gpu_data) + assert metrics["gpu_utilization_pct"] == 85.0 + assert metrics["temperature_c"] == 72.0 + assert metrics["power_draw_w"] == 200.5 + assert metrics["power_limit_w"] == 300.0 + assert metrics["vram_used_gb"] == round(8192 / 1024, 2) + assert metrics["vram_total_gb"] == round(16384 / 1024, 2) + assert metrics["vram_utilization_pct"] is not None + assert metrics["power_utilization_pct"] is not None + + def test_amd_primary_gpu_with_mock(self): + """get_primary_gpu_utilization returns correct dict with mocked amd-smi.""" + amd_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" + ) + _amd_spec = importlib.util.spec_from_file_location("test_amd2", amd_path) + assert _amd_spec is not None and _amd_spec.loader is not None + amd_mod = importlib.util.module_from_spec(_amd_spec) + + sys.modules["loggers"] = MagicMock() + sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock()) + + try: + _amd_spec.loader.exec_module(amd_mod) + except Exception: + pytest.skip("Could not load amd module") + + mock_json = json.dumps([{ + "usage": {"gfx_activity": "50"}, + "temperature": {"edge": "65"}, + "power": {"current_socket_power": "150", "power_cap": "250"}, + "vram": {"vram_used": 4096, "vram_total": 16384}, + }]) + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = mock_json + + with patch.object(subprocess, "run", return_value=mock_result): + result = amd_mod.get_primary_gpu_utilization() + assert result["available"] is True + assert result["gpu_utilization_pct"] == 50.0 + assert result["temperature_c"] == 65.0 + + def test_amd_smi_not_found_returns_unavailable(self): + """get_primary_gpu_utilization returns available=False when amd-smi is missing.""" + amd_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" + ) + _amd_spec = importlib.util.spec_from_file_location("test_amd3", amd_path) + assert _amd_spec is not None and _amd_spec.loader is not None + amd_mod = importlib.util.module_from_spec(_amd_spec) + + sys.modules["loggers"] = MagicMock() + sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock()) + + try: + _amd_spec.loader.exec_module(amd_mod) + except Exception: + pytest.skip("Could not load amd module") + + with patch.object( + subprocess, "run", side_effect=OSError("amd-smi not found") + ): + result = amd_mod.get_primary_gpu_utilization() + assert result["available"] is False + + def test_amd_timeout_returns_unavailable(self): + """get_primary_gpu_utilization handles timeout gracefully.""" + amd_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" + ) + _amd_spec = importlib.util.spec_from_file_location("test_amd4", amd_path) + assert _amd_spec is not None and _amd_spec.loader is not None + amd_mod = importlib.util.module_from_spec(_amd_spec) + + sys.modules["loggers"] = MagicMock() + sys.modules["loggers"].get_logger = MagicMock(return_value=MagicMock()) + + try: + _amd_spec.loader.exec_module(amd_mod) + except Exception: + pytest.skip("Could not load amd module") + + with patch.object( + subprocess, + "run", + side_effect=subprocess.TimeoutExpired("amd-smi", 5), + ): + result = amd_mod.get_primary_gpu_utilization() + assert result["available"] is False + + +# ============================================================================= +# TEST: hardware.py -- IS_ROCM branching to amd.py +# ============================================================================= + + +class TestHardwareAmdBranching: + """Verify hardware.py branches to amd.py when IS_ROCM is True.""" + + def test_hardware_imports_amd_module(self): + """hardware.py should import from amd module when IS_ROCM.""" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text() + assert "from . import amd" in source + + def test_hardware_branches_on_is_rocm_for_utilization(self): + """get_gpu_utilization should check IS_ROCM before choosing backend.""" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text() + # Find the get_gpu_utilization function + func_start = source.find("def get_gpu_utilization") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert "IS_ROCM" in func_body + assert "amd.get_primary_gpu_utilization" in func_body + + def test_hardware_branches_on_is_rocm_for_visible(self): + """get_visible_gpu_utilization should check IS_ROCM.""" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text() + func_start = source.find("def get_visible_gpu_utilization") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert "IS_ROCM" in func_body + assert "amd.get_visible_gpu_utilization" in func_body + + def test_hardware_branches_on_is_rocm_for_physical_count(self): + """get_physical_gpu_count should try amd.py when IS_ROCM.""" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text() + func_start = source.find("def get_physical_gpu_count") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert "IS_ROCM" in func_body + assert "amd.get_physical_gpu_count" in func_body + + +# ============================================================================= +# TEST: install_python_stack.py -- Windows AMD warning +# ============================================================================= + + +class TestWindowsRocmWarning: + """Verify Windows AMD GPU detection and warning message.""" + + def test_windows_amd_warning_in_source(self): + """install_python_stack.py should warn Windows AMD users.""" + source = _STACK_PATH.read_text() + assert "AMD GPU detected on Windows" in source + + def test_windows_amd_warning_checks_hipinfo_or_amdsmi(self): + """Warning should check for hipinfo or amd-smi.""" + source = _STACK_PATH.read_text() + assert "hipinfo" in source + assert "amd-smi" in source + + def test_windows_amd_warning_has_docs_link(self): + """Warning should include AMD docs link.""" + source = _STACK_PATH.read_text() + assert "docs.unsloth.ai/get-started/install-and-update/amd" in source + + +# ============================================================================= +# TEST: unsloth/kernels/utils.py -- is_rdna() expansion +# ============================================================================= + + +class TestIsRdnaExpansion: + """Verify is_rdna() covers RDNA2, RDNA3, RDNA3.5, RDNA4 architectures.""" + + def test_is_rdna_source_has_rdna2(self): + """is_rdna() should include RDNA2 architectures.""" + utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" + source = utils_path.read_text() + func_start = source.find("def is_rdna()") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert "gfx1030" in func_body + assert "gfx1031" in func_body + assert "gfx1032" in func_body + + def test_is_rdna_source_has_rdna3(self): + """is_rdna() should include RDNA3 architectures.""" + utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" + source = utils_path.read_text() + func_start = source.find("def is_rdna()") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert "gfx1100" in func_body + assert "gfx1101" in func_body + assert "gfx1102" in func_body + assert "gfx1103" in func_body + + def test_is_rdna_source_has_rdna35(self): + """is_rdna() should include RDNA3.5 architectures.""" + utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" + source = utils_path.read_text() + func_start = source.find("def is_rdna()") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert "gfx1150" in func_body + assert "gfx1151" in func_body + assert "gfx1152" in func_body + + def test_is_rdna_source_has_rdna4(self): + """is_rdna() should include RDNA4 architectures.""" + utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" + source = utils_path.read_text() + func_start = source.find("def is_rdna()") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert "gfx1200" in func_body + assert "gfx1201" in func_body + + def test_is_cdna_not_changed(self): + """is_cdna() should remain unchanged (no RDNA architectures added).""" + utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" + source = utils_path.read_text() + func_start = source.find("def is_cdna()") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert "gfx940" in func_body + assert "gfx941" in func_body + assert "gfx942" in func_body + assert "gfx950" in func_body + # RDNA architectures should NOT be in is_cdna + assert "gfx1030" not in func_body + assert "gfx1100" not in func_body + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 90f2d5d238..049afac96d 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -88,10 +88,22 @@ def is_cdna(): @functools.lru_cache(1) def is_rdna(): - """Detect ROCm-supported RDNA consumer/workstation GPUs (RDNA3, RDNA4).""" + """Detect ROCm-supported RDNA consumer/workstation GPUs (RDNA2, RDNA3, RDNA3.5, RDNA4).""" return is_hip() and triton.runtime.driver.active.get_current_target().arch in ( + # RDNA2 (Navi 21-24) + "gfx1030", + "gfx1031", + "gfx1032", + # RDNA3 (Navi 31-33) "gfx1100", "gfx1101", + "gfx1102", + "gfx1103", + # RDNA3.5 (Strix Point / Strix Halo) + "gfx1150", + "gfx1151", + "gfx1152", + # RDNA4 (Navi 48-44) "gfx1200", "gfx1201", )