diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index 6bd8600be3..747c515d68 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -65,6 +65,40 @@ def _parse_numeric(value: Any) -> Optional[float]: return None +def _parse_memory_mb(value: Any) -> Optional[float]: + """Parse a memory value from amd-smi output and return MB. + + Handles bare numbers (assumed MB), dict-shaped values with units + ({"value": 192, "unit": "GiB"}), and byte-scale heuristic fallback. + """ + unit = "" + raw_value = value + + if isinstance(value, dict): + unit = str(value.get("unit", "")).strip().lower() + raw_value = value.get("value") + + num = _parse_numeric(raw_value if isinstance(value, dict) else value) + if num is None: + return None + + # Explicit unit conversion + if "gib" in unit or "gb" in unit: + return num * 1024 + if "mib" in unit or "mb" in unit: + return num + if "kib" in unit or "kb" in unit: + return num / 1024 + if unit and ("b" in unit and "g" not in unit and "m" not in unit and "k" not in unit): + # Plain bytes + return num / (1024 * 1024) + + # No explicit unit -- heuristic: values > 10M are likely bytes + if num > 10_000_000: + return num / (1024 * 1024) + return num # Assume MB + + 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 @@ -107,34 +141,19 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: power_draw = None power_limit = None - # VRAM + # VRAM -- unit-aware parsing to handle varying amd-smi output formats. + # Newer amd-smi versions may return {"value": 192, "unit": "GiB"}. vram_data = gpu_data.get("vram", gpu_data.get("fb_memory_usage", {})) if isinstance(vram_data, dict): - vram_used_bytes = _parse_numeric( + vram_used_mb = _parse_memory_mb( vram_data.get("vram_used", vram_data.get("used")) ) - vram_total_bytes = _parse_numeric( + vram_total_mb = _parse_memory_mb( 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 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 > 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 > 10_000_000: # Likely bytes (>10M) - vram_total_mb = vram_total_bytes / (1024 * 1024) - else: # Likely already MB - vram_total_mb = vram_total_bytes + vram_used_mb = None + vram_total_mb = None # 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 diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 049e250237..4e9214deda 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -43,7 +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 +IS_ROCM: bool = False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py # ========== Detection ========== @@ -567,7 +567,17 @@ _visible_gpu_count: Optional[int] = None def _get_parent_visible_gpu_spec() -> Dict[str, Any]: - cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + # ROCm uses HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES in addition to + # CUDA_VISIBLE_DEVICES (which HIP also respects). Check ROCm-specific + # env vars first so multi-GPU AMD setups are handled correctly. + cuda_visible = None + if IS_ROCM: + cuda_visible = ( + os.environ.get("HIP_VISIBLE_DEVICES") + or os.environ.get("ROCR_VISIBLE_DEVICES") + ) + if cuda_visible is None: + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") if cuda_visible is None: return { diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 2ea84eede5..da2f47a04a 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -1450,7 +1450,7 @@ def detect_host() -> HostInfo: has_rocm = True break elif is_windows: - # Windows: validate actual AMD GPU presence (not just tool/DLL existence) + # Windows: prefer active probes that validate GPU presence for _cmd, _marker in ( (["hipinfo"], "gcnarchname"), (["amd-smi", "list"], "gpu"), @@ -1466,6 +1466,13 @@ def detect_host() -> HostInfo: if _marker in _result.stdout.lower(): has_rocm = True break + # Fallback: HIP runtime DLL indicates a working HIP installation + if not has_rocm and any( + Path(d).joinpath("amdhip64.dll").exists() + for d in os.environ.get("PATH", "").split(os.pathsep) + if d + ): + has_rocm = True return HostInfo( system = system, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index ac0d70f726..e884134a73 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -726,7 +726,29 @@ 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 and not _has_usable_nvidia_gpu(): - if shutil.which("hipinfo") or shutil.which("amd-smi"): + # Validate actual AMD GPU presence (not just tool existence) + _win_amd_gpu = False + for _wcmd, _wmarker in ( + (["hipinfo"], "gcnarchname"), + (["amd-smi", "list"], "gpu"), + ): + _wexe = shutil.which(_wcmd[0]) + if not _wexe: + continue + try: + _wr = subprocess.run( + [_wexe, *_wcmd[1:]], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + timeout=10, + ) + except Exception: + continue + if _wr.returncode == 0 and _wmarker in _wr.stdout.lower(): + _win_amd_gpu = True + break + if _win_amd_gpu: _safe_print( _dim(" Note:"), "AMD GPU detected on Windows. ROCm-enabled PyTorch must be",