From 7d722faedc3d2df60da982c84e072264caa0e457 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Apr 2026 21:22:45 +0000 Subject: [PATCH] Fix Windows ROCm install: idempotency guard, nvidia-smi fallback, Python version gates - Add torch.version.hip probe to _ensure_rocm_torch_windows() so steps 2b and 13 skip the 2.1-3.9 GB reinstall when ROCm torch is already healthy. Mirrors the existing Linux idempotency guard. - Add Windows fallback paths for nvidia-smi.exe in _has_usable_nvidia_gpu() (NVSMI dir + System32) matching the PowerShell scripts, so NVIDIA wins on mixed GPU systems even when nvidia-smi is not on PATH. - install.ps1: only narrow Python to 3.12 when AMD + torch (not --no-torch); set $PythonVersion = "3.12" so the existing winget auto-install flow works on clean AMD machines instead of hard-stopping. - install.ps1: handle empty $venvPyVer as an error in the AMD version check instead of silently continuing into a guaranteed pip failure. - setup.ps1: upgrade Python 3.12 version check from warning to hard error (exit 1), preventing multi-GB downloads that pip will reject on non-3.12. - install_python_stack.py: use os.environ["ProgramFiles"] for ROCm root fallback instead of hardcoded C:\Program Files. - Remove unused _HIP_SDK_DOWNLOAD_URL constant. --- install.ps1 | 22 ++++++++--------- studio/install_python_stack.py | 44 ++++++++++++++++++++++++++++++---- studio/setup.ps1 | 19 +++++++++------ 3 files changed, 62 insertions(+), 23 deletions(-) diff --git a/install.ps1 b/install.ps1 index a8957520ed..c1dbb2ddbd 100644 --- a/install.ps1 +++ b/install.ps1 @@ -784,11 +784,12 @@ shell.Run cmd, 0, False # ── Install Python if no compatible version found ── # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. - # AMD path requires Python 3.12 specifically (Radeon only publishes cp312 - # wheels for Windows), so we pass a narrower preference list AND skip the - # winget auto-install (the user must have 3.12 already). - if ($HasAmdGpu) { + # AMD + torch path requires Python 3.12 specifically (Radeon only publishes + # cp312 wheels for Windows). AMD + --no-torch (GGUF-only) is fine with any + # 3.11-3.13, since the cp312 constraint only applies to the ROCm wheels. + if ($HasAmdGpu -and -not $SkipTorch) { $PythonPreferred = @("3.12") + $PythonVersion = "3.12" } else { $PythonPreferred = @("3.13", "3.12", "3.11") } @@ -796,13 +797,6 @@ shell.Run cmd, 0, False if ($DetectedPython) { step "python" "Python $($DetectedPython.Version) already installed" } - if (-not $DetectedPython -and $HasAmdGpu) { - Write-Host "[ERROR] AMD ROCm path requires Python 3.12 (Radeon's Windows wheels are cp312 only)." -ForegroundColor Red - Write-Host " Install Python 3.12 from https://www.python.org/downloads/ or via:" -ForegroundColor Yellow - Write-Host " winget install -e --id Python.Python.3.12" -ForegroundColor Yellow - Write-Host " Then re-run this installer." -ForegroundColor Yellow - return - } if (-not $DetectedPython) { substep "installing Python ${PythonVersion}..." $pythonPackageId = "Python.Python.$PythonVersion" @@ -1063,7 +1057,11 @@ shell.Run cmd, 0, False try { $venvPyVer = (& $VenvPython -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')" 2>$null | Out-String).Trim() } catch {} - if ($venvPyVer -and $venvPyVer -ne "3.12") { + if (-not $venvPyVer) { + Write-Host "[ERROR] Could not determine venv Python version." -ForegroundColor Red + return + } + if ($venvPyVer -ne "3.12") { Write-Host "[ERROR] Radeon Windows ROCm wheels require Python 3.12 (venv has $venvPyVer)." -ForegroundColor Red Write-Host " Install Python 3.12 from https://www.python.org/downloads/ and re-run." -ForegroundColor Yellow return diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index aea1bc085a..52381bfc9a 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -164,9 +164,6 @@ _ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = { # developer toolkit is NOT required for running torch. _DEFAULT_WINDOWS_ROCM_VERSION: tuple[int, int] = (7, 2) _AMD_RADEON_DRIVER_URL = "https://www.amd.com/en/support/download/drivers.html" -_HIP_SDK_DOWNLOAD_URL = ( - "https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -) # bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix # (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every @@ -363,7 +360,9 @@ def _detect_rocm_version_windows() -> tuple[int, int] | None: if ver is not None: return ver - rocm_root = r"C:\Program Files\AMD\ROCm" + rocm_root = os.path.join( + os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm" + ) if os.path.isdir(rocm_root): best: tuple[int, int] | None = None try: @@ -426,6 +425,25 @@ def _has_rocm_gpu_windows() -> bool: 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 and IS_WINDOWS: + # nvidia-smi.exe is often absent from PATH on Windows even with a + # valid driver. Match the fallback paths used in install.ps1 / + # setup.ps1 so the NVIDIA-wins-on-mixed-systems rule is consistent + # between the PowerShell and Python install paths. + _candidates = [ + os.path.join( + os.environ.get("ProgramFiles", r"C:\Program Files"), + r"NVIDIA Corporation\NVSMI\nvidia-smi.exe", + ), + os.path.join( + os.environ.get("SystemRoot", r"C:\Windows"), + r"System32\nvidia-smi.exe", + ), + ] + for _c in _candidates: + if os.path.isfile(_c): + exe = _c + break if not exe: return False try: @@ -475,6 +493,24 @@ def _ensure_rocm_torch_windows() -> None: if not _has_rocm_gpu_windows(): return + # Skip when torch already links against ROCm -- mirrors the Linux + # has_hip_torch probe (line ~622) and makes this function idempotent. + # Without this guard, steps 2b and 13 in install_python_stack() would + # each re-download the full 2.1-3.9 GB wheel set even when the first + # call (or a prior setup.ps1 / install.ps1 run) already succeeded. + try: + _probe = subprocess.run( + [sys.executable, "-c", + "import torch; print(getattr(torch.version,'hip','') or '')"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 30, + ) + if _probe.returncode == 0 and _probe.stdout.decode().strip(): + return # already has ROCm torch + except Exception: + pass + # Radeon wheels are cp312 only. Warn (do not crash) when the venv's # Python is not 3.12 -- pip will fail anyway with a clearer message. if (sys.version_info.major, sys.version_info.minor) != (3, 12): diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 3500b3caf4..3461194e61 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1681,16 +1681,21 @@ if ($CuTag -eq "rocm") { exit 1 } - # Radeon wheels are cp312 only. Warn loudly when the venv's Python is - # a different minor version so the pip error makes sense. We do not - # exit here because setup.ps1 is also invoked as "unsloth studio update" - # inside a venv the user may have created manually. + # Radeon wheels are cp312 only. Hard-stop when the venv's Python is a + # different minor version -- continuing would download multi-GB wheels + # that pip will reject with a confusing incompatible-wheel error. try { $venvPyVer = (& python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')" 2>$null | Out-String).Trim() } catch { $venvPyVer = "" } - if ($venvPyVer -and $venvPyVer -ne "3.12") { - substep "Warning: Radeon Windows ROCm wheels require Python 3.12, venv has $venvPyVer." "Yellow" - substep "Re-create the venv with Python 3.12 from https://python.org if pip fails below." "Yellow" + if (-not $venvPyVer) { + Write-Host "[FAILED] Could not determine venv Python version." -ForegroundColor Red + Write-Host " Re-create the venv with Python 3.12 from https://python.org and re-run." -ForegroundColor Yellow + exit 1 + } + if ($venvPyVer -ne "3.12") { + Write-Host "[FAILED] Radeon Windows ROCm wheels require Python 3.12, venv has $venvPyVer." -ForegroundColor Red + Write-Host " Re-create the venv with Python 3.12 from https://python.org and re-run." -ForegroundColor Yellow + exit 1 } substep ("installing Radeon ROCm SDK + PyTorch for rocm-rel-{0}.{1}.x from repo.radeon.com..." -f $RocmReleaseVersion.Major, $RocmReleaseVersion.Minor)