From c6a9585659e56910eee40371d03667dbe331c3dd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 10 Apr 2026 15:14:11 +0000 Subject: [PATCH 01/15] Install ROCm PyTorch on Windows AMD via repo.radeon.com AMD support on Windows fell back to CPU-only torch because install.ps1, studio/setup.ps1, and the Windows branch of studio/install_python_stack.py only detected nvidia-smi. This fixes unslothai/unsloth#4280 by teaching the installers to pick ROCm torch from repo.radeon.com when an AMD GPU plus HIP SDK 7.1.x or 7.2.x is present. Changes - install.ps1: Get-HipSdkVersion + Get-RocmWheelUrls helpers, AMD GPU detection (WMI Win32_VideoController), Python 3.12 enforcement on the AMD path (Radeon wheels are cp312 only), and a dedicated AMD torch install branch that skips bitsandbytes. - studio/setup.ps1: mirrors the install.ps1 helpers (self-contained copy), adds an AMD branch to the torch install flow, and teaches the stale-venv check to match both +rocm and +rocmsdk suffixes so ROCm minor updates do not trigger spurious venv rebuilds. - studio/install_python_stack.py: new _ROCM_WINDOWS_TORCH_WHEELS mapping, _detect_rocm_version_windows (HIP_PATH primary + ProgramFiles scan fallback, uses ntpath so path parsing works on Linux test runners), _has_rocm_gpu_windows via PowerShell WMI, and a new _ensure_rocm_torch_windows helper that respects the NVIDIA-wins rule. The bnb install section returns early on Windows because there is no Windows ROCm wheel (bitsandbytes-foundation/bitsandbytes#1844). NVIDIA, CPU-only, Linux AMD, and macOS paths are untouched. On Windows NVIDIA+AMD mixed hosts NVIDIA takes precedence, matching install.sh behaviour. --- install.ps1 | 298 +++++++++++++++++++++++++++----- studio/install_python_stack.py | 299 ++++++++++++++++++++++++++------- studio/setup.ps1 | 177 ++++++++++++++++++- 3 files changed, 671 insertions(+), 103 deletions(-) diff --git a/install.ps1 b/install.ps1 index a2acd6c4ea..5d70be4b4d 100644 --- a/install.ps1 +++ b/install.ps1 @@ -199,6 +199,86 @@ function Install-UnslothStudio { } } + # ── AMD Windows ROCm helpers ── + # Detect AMD HIP SDK 7.1.x / 7.2.x on Windows. Returns a hashtable + # @{ Major = 7; Minor = 2; Path = "C:\Program Files\AMD\ROCm\7.2" } + # or $null when the SDK is missing, unreadable, or at an unsupported + # version. Callers differentiate "not installed" vs "installed but wrong + # version" by re-parsing $env:HIP_PATH themselves. + function Get-HipSdkVersion { + # Primary signal: HIP_PATH env var, e.g. "C:\Program Files\AMD\ROCm\7.2\" + # The final path component is the version. Validate bin\ exists so + # a stale env var pointing at a half-uninstalled SDK is rejected. + $parseComponent = { + param($name) + if ([string]::IsNullOrWhiteSpace($name)) { return $null } + $m = [regex]::Match($name.Trim(), '^(\d+)\.(\d+)') + if ($m.Success) { + return @{ Major = [int]$m.Groups[1].Value; Minor = [int]$m.Groups[2].Value } + } + return $null + } + + $hipPath = $env:HIP_PATH + if ($hipPath) { + $trimmed = $hipPath.TrimEnd('\','/') + if (Test-Path (Join-Path $trimmed 'bin')) { + $parsed = & $parseComponent (Split-Path $trimmed -Leaf) + if ($parsed) { + $parsed.Path = $trimmed + return $parsed + } + } + } + + # Fallback: scan C:\Program Files\AMD\ROCm\ + $rocmRoot = Join-Path $env:ProgramFiles 'AMD\ROCm' + if (Test-Path $rocmRoot) { + $best = $null + try { + $dirs = Get-ChildItem -Path $rocmRoot -Directory -ErrorAction SilentlyContinue + } catch { $dirs = @() } + foreach ($d in $dirs) { + if (-not (Test-Path (Join-Path $d.FullName 'bin'))) { continue } + $parsed = & $parseComponent $d.Name + if ($null -eq $parsed) { continue } + if ($null -eq $best -or + $parsed.Major -gt $best.Major -or + ($parsed.Major -eq $best.Major -and $parsed.Minor -gt $best.Minor)) { + $parsed.Path = $d.FullName + $best = $parsed + } + } + if ($best) { return $best } + } + + return $null + } + + # Map a detected HIP SDK version to Radeon's Windows torch wheels. + # Returns @{ Torch = ...; Torchvision = ...; Torchaudio = ... } or $null + # when the version is unsupported. Wheels are cp312 only. + function Get-RocmWheelUrls { + param([Parameter(Mandatory = $true)]$Version) + $base721 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/' + $base711 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/' + if ($Version.Major -eq 7 -and $Version.Minor -eq 2) { + return @{ + Torch = $base721 + 'torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + Torchvision = $base721 + 'torchvision-0.24.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + Torchaudio = $base721 + 'torchaudio-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + } + } + if ($Version.Major -eq 7 -and $Version.Minor -eq 1) { + return @{ + Torch = $base711 + 'torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + Torchvision = $base711 + 'torchvision-0.24.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + Torchaudio = $base711 + 'torchaudio-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + } + } + return $null + } + function New-StudioShortcuts { param( [Parameter(Mandatory = $true)][string]$UnslothExePath @@ -523,6 +603,72 @@ shell.Run cmd, 0, False return } + # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── + # Runs before Python detection so the AMD/ROCm path can require Python + # 3.12 (the only cp tag Radeon publishes Windows wheels for). + $HasNvidiaSmi = $false + $NvidiaSmiExe = $null + try { + $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue + if ($nvSmiCmd) { + & $nvSmiCmd.Source *> $null + if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source } + } + } catch {} + if (-not $HasNvidiaSmi) { + foreach ($p in @( + "$env:ProgramFiles\NVIDIA Corporation\NVSMI\nvidia-smi.exe", + "$env:SystemRoot\System32\nvidia-smi.exe" + )) { + if (Test-Path $p) { + try { + & $p *> $null + if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break } + } catch {} + } + } + } + + # AMD GPU presence via WMI -- always available on Windows, no elevation, + # no dependency on HIP SDK being installed. We deliberately avoid + # hipinfo.exe here because it ships inside the HIP SDK and would create + # a chicken-and-egg where we cannot prompt "install HIP SDK" on the hosts + # that need the prompt. + $HasAmdGpu = $false + try { + $videoControllers = Get-CimInstance Win32_VideoController -ErrorAction SilentlyContinue + if ($videoControllers) { + $amdMatches = @($videoControllers | Where-Object { $_.Name -match 'AMD|Radeon' }) + if ($amdMatches.Count -gt 0) { $HasAmdGpu = $true } + } + } catch {} + # NVIDIA wins on mixed systems -- matches Linux install.sh behaviour so + # users with both cards get the NVIDIA torch path they expect. + if ($HasNvidiaSmi) { $HasAmdGpu = $false } + + # Resolve HIP SDK when we are taking the AMD path. We still probe here + # even when $HasAmdGpu is $false so the status line is informative on + # mixed NVIDIA+AMD hosts. + $HipSdkVersion = $null + if ($HasAmdGpu) { + $HipSdkVersion = Get-HipSdkVersion + } + + if ($HasNvidiaSmi) { + step "gpu" "NVIDIA GPU detected" + } elseif ($HasAmdGpu) { + if ($HipSdkVersion) { + step "gpu" ("AMD GPU detected (HIP SDK {0}.{1})" -f $HipSdkVersion.Major, $HipSdkVersion.Minor) + } else { + step "gpu" "AMD GPU detected (HIP SDK missing)" "Yellow" + substep "Install HIP SDK from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" "Yellow" + substep "and re-run this installer." "Yellow" + } + } else { + step "gpu" "none (chat-only / GGUF)" "Yellow" + substep "Training and GPU inference require an NVIDIA or AMD GPU with drivers installed." "Yellow" + } + # ── Helper: detect a working Python 3.11-3.13 on the system ── # Returns the version string (e.g. "3.13") or "" if none found. # Uses try-catch + stderr redirection so that App Execution Alias stubs @@ -551,15 +697,31 @@ shell.Run cmd, 0, False # Returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. # The resolved Path is passed to `uv venv --python` to prevent uv from # re-resolving the version string back to a conda interpreter. + # + # $PreferredVersions is a list of minor versions to search, in priority + # order. Defaults to 3.13, 3.12, 3.11. AMD/ROCm callers pass @("3.12") + # because Radeon's Windows wheels are cp312 only. function Find-CompatiblePython { + param([string[]]$PreferredVersions = @("3.13", "3.12", "3.11")) + + # Build a version regex from the requested minor list. We accept + # "3.." where is in the preferred list. + $minors = @() + foreach ($v in $PreferredVersions) { + if ($v -match '^3\.(\d+)$') { $minors += $Matches[1] } + } + if ($minors.Count -eq 0) { return $null } + $minorAlt = ($minors | Sort-Object -Unique) -join '|' + $verRegex = "Python (3\.(?:$minorAlt))\.\d+" + # Try the Python Launcher first (most reliable on Windows) # py.exe resolves to the standard CPython install, not conda. $pyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue if ($pyLauncher -and $pyLauncher.Source -notmatch $script:CondaSkipPattern) { - foreach ($minor in @("3.13", "3.12", "3.11")) { + foreach ($minor in $PreferredVersions) { try { $out = & $pyLauncher.Source "-$minor" --version 2>&1 | Out-String - if ($out -match "Python (3\.1[1-3])\.\d+") { + if ($out -match $verRegex) { $ver = $Matches[1] # Resolve the actual executable path and verify it is not conda-based $resolvedExe = (& $pyLauncher.Source "-$minor" -c "import sys; print(sys.executable)" 2>$null | Out-String).Trim() @@ -584,7 +746,7 @@ shell.Run cmd, 0, False if (Test-IsCondaPython $cmd.Source) { continue } try { $out = & $cmd.Source --version 2>&1 | Out-String - if ($out -match "Python (3\.1[1-3])\.\d+") { + if ($out -match $verRegex) { return @{ Version = $Matches[1]; Path = $cmd.Source } } } catch {} @@ -593,12 +755,27 @@ shell.Run cmd, 0, False return $null } - # ── Install Python if no compatible version (3.11-3.13) found ── + # ── Install Python if no compatible version found ── # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. - $DetectedPython = Find-CompatiblePython + # 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) { + $PythonPreferred = @("3.12") + } else { + $PythonPreferred = @("3.13", "3.12", "3.11") + } + $DetectedPython = Find-CompatiblePython -PreferredVersions $PythonPreferred 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" @@ -615,7 +792,7 @@ shell.Run cmd, 0, False Refresh-SessionPath # Re-detect after install (PATH may have changed) - $DetectedPython = Find-CompatiblePython + $DetectedPython = Find-CompatiblePython -PreferredVersions $PythonPreferred if (-not $DetectedPython) { # Python still not functional after winget -- force reinstall. @@ -630,7 +807,7 @@ shell.Run cmd, 0, False } catch { $wingetExit = 1 } $ErrorActionPreference = $prevEAP Refresh-SessionPath - $DetectedPython = Find-CompatiblePython + $DetectedPython = Find-CompatiblePython -PreferredVersions $PythonPreferred } if (-not $DetectedPython) { @@ -721,35 +898,9 @@ shell.Run cmd, 0, False substep "$VenvDir" } - # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── - $HasNvidiaSmi = $false - $NvidiaSmiExe = $null - try { - $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue - if ($nvSmiCmd) { - & $nvSmiCmd.Source *> $null - if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source } - } - } catch {} - if (-not $HasNvidiaSmi) { - foreach ($p in @( - "$env:ProgramFiles\NVIDIA Corporation\NVSMI\nvidia-smi.exe", - "$env:SystemRoot\System32\nvidia-smi.exe" - )) { - if (Test-Path $p) { - try { - & $p *> $null - if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break } - } catch {} - } - } - } - if ($HasNvidiaSmi) { - step "gpu" "NVIDIA GPU detected" - } else { - step "gpu" "none (chat-only / GGUF)" "Yellow" - substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" - } + # GPU detection moved earlier in the flow (see "Detect GPU" block above + # the Python detection); $HasNvidiaSmi, $NvidiaSmiExe, $HasAmdGpu, and + # $HipSdkVersion are already populated by this point. # ── Choose the correct PyTorch index URL based on driver CUDA version ── # Mirrors Get-PytorchCudaTag in setup.ps1. @@ -771,10 +922,22 @@ shell.Run cmd, 0, False substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow" return "$baseUrl/cu126" } - $TorchIndexUrl = Get-TorchIndexUrl + + # The AMD/ROCm path does not use a --index-url; instead it installs + # explicit wheel URLs from repo.radeon.com. $TorchIndexUrl stays $null + # on that branch so the NVIDIA/CPU path is visibly bypassed. + $TorchIndexUrl = $null + $RocmWheelUrls = $null + if ($HasAmdGpu) { + if ($HipSdkVersion) { + $RocmWheelUrls = Get-RocmWheelUrls -Version $HipSdkVersion + } + } else { + $TorchIndexUrl = Get-TorchIndexUrl + } # ── Print CPU-only hint when no GPU detected ── - if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") { + if (-not $SkipTorch -and -not $HasAmdGpu -and $TorchIndexUrl -like "*/cpu") { Write-Host "" substep "No NVIDIA GPU detected." "Yellow" substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow" @@ -833,6 +996,65 @@ shell.Run cmd, 0, False Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return } + if ($StudioLocalInstall) { + substep "overlaying local repo (editable)..." + $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + if ($overlayExit -ne 0) { + Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red + return + } + } + } elseif ($HasAmdGpu) { + if ($SkipTorch) { + substep "skipping PyTorch (--no-torch flag set)." "Yellow" + } elseif (-not $RocmWheelUrls) { + if (-not $HipSdkVersion) { + Write-Host "[ERROR] AMD GPU detected but HIP SDK is not installed." -ForegroundColor Red + Write-Host " Download it from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -ForegroundColor Yellow + Write-Host " and re-run this installer." -ForegroundColor Yellow + } else { + Write-Host "[ERROR] AMD HIP SDK $($HipSdkVersion.Major).$($HipSdkVersion.Minor) detected." -ForegroundColor Red + Write-Host " Unsloth requires HIP SDK 7.1 or later on Windows. Please update from" -ForegroundColor Yellow + Write-Host " https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -ForegroundColor Yellow + } + return + } else { + substep "installing PyTorch ROCm wheels from repo.radeon.com..." + substep "(first wheel is ~780 MB; download may take a few minutes)" + $torchInstallExit = Invoke-InstallCommand { + uv pip install --python $VenvPython ` + $RocmWheelUrls.Torch $RocmWheelUrls.Torchvision $RocmWheelUrls.Torchaudio + } + if ($torchInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red + Write-Host " Verify HIP SDK $($HipSdkVersion.Major).$($HipSdkVersion.Minor) is installed and wheels are reachable." -ForegroundColor Yellow + return + } + } + + substep "installing unsloth (this may take a few minutes)..." + # --no-deps prevents uv from re-resolving torch back to default PyPI + # (which would strip the +rocm suffix). bitsandbytes has no Windows + # ROCm wheel so it is deliberately NOT installed here; 4-bit + # quantization is not available on Windows AMD yet. + if ($SkipTorch) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.4.4" unsloth-zoo } + if ($baseInstallExit -eq 0) { + $NoTorchReq = Find-NoTorchRuntimeFile + if ($NoTorchReq) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + } + } + } elseif ($StudioLocalInstall) { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.4.4" unsloth-zoo } + } else { + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" } + } + if ($baseInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red + return + } + if ($StudioLocalInstall) { substep "overlaying local repo (editable)..." $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a046f8f892..588671ab11 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -43,6 +43,46 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { } _PYTORCH_WHL_BASE = "https://download.pytorch.org/whl" +# Windows AMD ROCm torch wheels live at repo.radeon.com, not download.pytorch.org. +# Keyed by HIP SDK (major, minor). Wheels are cp312 only and require the HIP SDK +# to be pre-installed. Older releases (rocm-rel-6.4.4) use a nested layout and +# alpha version strings that are fragile to match on, so only 7.1.x / 7.2.x are +# supported here -- older HIP SDKs fall through with a pointer to the download +# page. +_ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = { + (7, 2): { + "torch": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" + "torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl" + ), + "torchvision": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" + "torchvision-0.24.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl" + ), + "torchaudio": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" + "torchaudio-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl" + ), + }, + (7, 1): { + "torch": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/" + "torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl" + ), + "torchvision": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/" + "torchvision-0.24.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl" + ), + "torchaudio": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/" + "torchaudio-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl" + ), + }, +} +_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 # AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI. @@ -200,6 +240,104 @@ def _has_rocm_gpu() -> bool: return False +def _detect_rocm_version_windows() -> tuple[int, int] | None: + """Return (major, minor) of the installed HIP SDK on Windows, or None. + + Primary signal: HIP_PATH env var set by the HIP SDK installer, e.g. + ``C:\\Program Files\\AMD\\ROCm\\7.2\\``. The trailing component is the + version. We also sanity-check that ``\\bin`` exists so we do + not trust a stale env var pointing at a half-uninstalled SDK. + + Fallback: scan ``C:\\Program Files\\AMD\\ROCm\\`` and return the highest + numeric subdirectory. + """ + if not IS_WINDOWS: + return None + # Use ntpath explicitly rather than os.path so the logic is also + # testable on Linux runners where os.path is posixpath (basename of + # "C:\\foo\\7.2" returns the whole string on posix, breaking version + # extraction even though real Windows handles it correctly). + import ntpath as _ntpath + import re as _re + + def _parse_component(name: str) -> tuple[int, int] | None: + # Accept "7.2" and "7.2.1", ignore anything else. + m = _re.match(r"^(\d+)\.(\d+)", name.strip()) + if not m: + return None + return int(m.group(1)), int(m.group(2)) + + hip_path = os.environ.get("HIP_PATH", "").strip() + if hip_path: + trimmed = hip_path.rstrip("\\/") + bin_dir = _ntpath.join(trimmed, "bin") + # Only trust HIP_PATH when the bin folder is actually present. A + # broken uninstall can leave the env var pointing at a ghost dir. + if os.path.isdir(bin_dir): + ver = _parse_component(_ntpath.basename(trimmed)) + if ver is not None: + return ver + + rocm_root = r"C:\Program Files\AMD\ROCm" + if os.path.isdir(rocm_root): + best: tuple[int, int] | None = None + try: + entries = os.listdir(rocm_root) + except OSError: + entries = [] + for entry in entries: + sub = _ntpath.join(rocm_root, entry) + if not os.path.isdir(sub): + continue + ver = _parse_component(entry) + if ver is None: + continue + if os.path.isdir(_ntpath.join(sub, "bin")): + if best is None or ver > best: + best = ver + if best is not None: + return best + + return None + + +def _has_rocm_gpu_windows() -> bool: + """Return True when a Radeon/AMD GPU is visible in WMI Win32_VideoController. + + We deliberately avoid ``hipinfo.exe`` here because it lives inside the + HIP SDK -- if we used it to decide whether to prompt the user to install + the HIP SDK we would never trigger the prompt on the hosts that need it + most. WMI is always available on Windows and needs no elevation. + """ + if not IS_WINDOWS: + return False + ps_cmd = ( + "Get-CimInstance Win32_VideoController -ErrorAction SilentlyContinue " + "| Where-Object { $_.Name -match 'AMD|Radeon' } " + "| Measure-Object | Select-Object -ExpandProperty Count" + ) + for exe_name in ("pwsh", "powershell"): + exe = shutil.which(exe_name) + if not exe: + continue + try: + result = subprocess.run( + [exe, "-NoProfile", "-NonInteractive", "-Command", ps_cmd], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 15, + ) + except Exception: + continue + if result.returncode != 0: + continue + raw = (result.stdout or "").strip() + if raw.isdigit() and int(raw) > 0: + return True + 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") @@ -218,25 +356,101 @@ def _has_usable_nvidia_gpu() -> bool: return result.returncode == 0 and "GPU " in result.stdout +def _ensure_rocm_torch_windows() -> None: + """Install Radeon's Windows ROCm torch wheels when an AMD GPU + HIP SDK + are both present. Called from _ensure_rocm_torch(). + + Silently returns when no AMD GPU is visible, so NVIDIA and CPU-only + Windows hosts are never touched. When an AMD GPU is present but the + HIP SDK is missing or too old, prints a pointer to the HIP SDK + download page and returns without raising -- the Linux helper has the + same shape. NVIDIA takes precedence on mixed AMD+NVIDIA hosts so + install.ps1 and setup.ps1 (which install CUDA torch in that case) + are not clobbered. + """ + # NVIDIA wins on mixed hosts -- matches the Linux branch and avoids + # overwriting a freshly installed CUDA torch with ROCm wheels. + if _has_usable_nvidia_gpu(): + return + if not _has_rocm_gpu_windows(): + return + + ver = _detect_rocm_version_windows() + if ver is None: + _safe_print( + _red( + " AMD GPU detected but HIP SDK was not found. Install it " + f"from {_HIP_SDK_DOWNLOAD_URL} and re-run setup." + ) + ) + return + + wheels = _ROCM_WINDOWS_TORCH_WHEELS.get(ver) + if wheels is None: + _safe_print( + _red( + f" HIP SDK {ver[0]}.{ver[1]} detected. Unsloth on Windows " + f"requires HIP SDK 7.1 or 7.2. Please update from " + f"{_HIP_SDK_DOWNLOAD_URL}" + ) + ) + return + + # 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): + _safe_print( + _red( + f" Radeon Windows ROCm wheels require Python 3.12. " + f"Found Python {sys.version_info.major}.{sys.version_info.minor}. " + f"Install Python 3.12 from https://python.org and re-run." + ) + ) + return + + _safe_print( + _dim( + f" HIP SDK {ver[0]}.{ver[1]} -- installing torch from " + f"repo.radeon.com/rocm/windows/" + ) + ) + pip_install( + f"ROCm torch (Windows, HIP SDK {ver[0]}.{ver[1]})", + "--force-reinstall", + "--no-cache-dir", + wheels["torch"], + wheels["torchvision"], + wheels["torchaudio"], + constrain = False, + ) + + def _ensure_rocm_torch() -> None: """Reinstall torch with ROCm wheels when the venv received CPU-only torch. - Runs only on Linux x86_64 hosts where an AMD GPU is present and the - ROCm runtime is detectable (rocminfo / amd-smi / hipconfig / - rocm-core package). No-op when torch already links against HIP - (ROCm), on Windows / macOS, on non-x86_64 Linux (PyTorch does not - publish ROCm wheels for aarch64 / arm64), or on mixed AMD+NVIDIA - hosts (NVIDIA takes precedence). + Linux x86_64: downloads ROCm wheels from download.pytorch.org for the + detected /opt/rocm version. Windows x86_64: downloads ROCm wheels from + repo.radeon.com for the detected HIP SDK version (requires Python 3.12 + because Radeon only publishes cp312 wheels). No-op on macOS, on + non-x86_64 hosts (PyTorch does not publish ROCm wheels for aarch64 / + arm64), or on mixed AMD+NVIDIA hosts (NVIDIA takes precedence). No-op + when torch already links against HIP on Linux. + Uses pip_install() to respect uv, constraints, and --python targeting. """ # Explicit OS / architecture guards so the helper is safe to call - # from any context -- PyTorch only publishes ROCm wheels for - # linux_x86_64, so aarch64 / arm64 hosts must skip this repair path - # instead of failing the update with a missing-wheel error. - if IS_WINDOWS or IS_MACOS: + # from any context -- PyTorch only publishes ROCm wheels for x86_64, + # so aarch64 / arm64 hosts must skip this repair path instead of + # failing the update with a missing-wheel error. + if IS_MACOS: return if platform.machine().lower() not in {"x86_64", "amd64"}: return + + if IS_WINDOWS: + _ensure_rocm_torch_windows() + return + # NVIDIA takes precedence on mixed hosts -- but only if an actual GPU is usable if _has_usable_nvidia_gpu(): return @@ -314,6 +528,12 @@ def _ensure_rocm_torch() -> None: # continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls # back to PyPI when the pre-release URL is unreachable. if rocm_torch_ready: + # bitsandbytes has no official Windows ROCm wheel + # (bitsandbytes-foundation/bitsandbytes#1844), so skip it entirely + # on Windows AMD. 4-bit quantization is not available on that path + # yet -- callers should install 16-bit or use GGUF inference. + if IS_WINDOWS: + return _bnb_url = _bnb_rocm_prerelease_url() _bnb_installed = False if _bnb_url is not None: @@ -762,9 +982,9 @@ def install_python_stack() -> int: base_total = 10 if IS_WINDOWS else 11 if IS_MACOS: base_total -= 1 # triton step is skipped on macOS - # ROCm torch check steps (Linux only, non-macOS, non-no-torch): + # ROCm torch check steps (Linux + Windows, non-macOS, non-no-torch): # one early check (step 2b) and one final repair (step 13). - if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: + if not IS_MACOS and not NO_TORCH: base_total += 2 _TOTAL = (base_total - 1) if skip_base else base_total @@ -892,50 +1112,12 @@ def install_python_stack() -> int: # 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the # venv received CPU-only torch (common when pip resolves torch from PyPI). # Must come immediately after base packages so torch is present for inspection. - if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: + # On Linux this pulls from download.pytorch.org; on Windows from + # repo.radeon.com. _ensure_rocm_torch() dispatches internally. + if not IS_MACOS and not NO_TORCH: _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 and not _has_usable_nvidia_gpu(): - # Validate actual AMD GPU presence (not just tool existence) - import re as _re_win - - def _win_amd_smi_has_gpu(stdout: str) -> bool: - return bool(_re_win.search(r"(?im)^gpu\s*[:\[]\s*\d", stdout)) - - _win_amd_gpu = False - for _wcmd, _check_fn in ( - (["hipinfo"], lambda out: "gcnarchname" in out.lower()), - (["amd-smi", "list"], _win_amd_smi_has_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 _check_fn(_wr.stdout): - _win_amd_gpu = True - break - if _win_amd_gpu: - _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( @@ -1050,11 +1232,12 @@ def install_python_stack() -> int: ) # 13. AMD ROCm: final torch repair. Multiple install steps above can - # pull in CUDA torch from PyPI (base packages, extras, overrides, - # studio deps, etc.). Running the repair as the very last step - # ensures ROCm torch is in place at runtime, regardless of which - # intermediate step clobbered it. - if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: + # pull in CUDA / CPU torch from PyPI (base packages, extras, + # overrides, studio deps, etc.). Running the repair as the very + # last step ensures ROCm torch is in place at runtime, regardless + # of which intermediate step clobbered it. Same behavior on Linux + # (download.pytorch.org wheels) and Windows (repo.radeon.com). + if not IS_MACOS and not NO_TORCH: _progress("ROCm torch (final)") _ensure_rocm_torch() diff --git a/studio/setup.ps1 b/studio/setup.ps1 index c3a8cd71ca..a74b116ff5 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -266,6 +266,83 @@ function Get-PytorchCudaTag { return "cu126" } +# ───────────────────────────────────────────── +# AMD Windows ROCm helpers (self-contained copy -- kept in sync with +# install.ps1; do not dot-source install.ps1 because setup.ps1 is also +# invoked standalone from "unsloth studio update"). +# ───────────────────────────────────────────── + +# Detect AMD HIP SDK 7.1.x / 7.2.x on Windows. Returns +# @{ Major = 7; Minor = 2; Path = "C:\Program Files\AMD\ROCm\7.2" } +# or $null when not installed / unreadable. +function Get-HipSdkVersion { + $parseComponent = { + param($name) + if ([string]::IsNullOrWhiteSpace($name)) { return $null } + $m = [regex]::Match($name.Trim(), '^(\d+)\.(\d+)') + if ($m.Success) { + return @{ Major = [int]$m.Groups[1].Value; Minor = [int]$m.Groups[2].Value } + } + return $null + } + + $hipPath = $env:HIP_PATH + if ($hipPath) { + $trimmed = $hipPath.TrimEnd('\','/') + if (Test-Path (Join-Path $trimmed 'bin')) { + $parsed = & $parseComponent (Split-Path $trimmed -Leaf) + if ($parsed) { + $parsed.Path = $trimmed + return $parsed + } + } + } + + $rocmRoot = Join-Path $env:ProgramFiles 'AMD\ROCm' + if (Test-Path $rocmRoot) { + $best = $null + try { + $dirs = Get-ChildItem -Path $rocmRoot -Directory -ErrorAction SilentlyContinue + } catch { $dirs = @() } + foreach ($d in $dirs) { + if (-not (Test-Path (Join-Path $d.FullName 'bin'))) { continue } + $parsed = & $parseComponent $d.Name + if ($null -eq $parsed) { continue } + if ($null -eq $best -or + $parsed.Major -gt $best.Major -or + ($parsed.Major -eq $best.Major -and $parsed.Minor -gt $best.Minor)) { + $parsed.Path = $d.FullName + $best = $parsed + } + } + if ($best) { return $best } + } + return $null +} + +# Map a detected HIP SDK version to Radeon's Windows torch wheels. +# Returns @{ Torch = ...; Torchvision = ...; Torchaudio = ... } or $null. +function Get-RocmWheelUrls { + param([Parameter(Mandatory = $true)]$Version) + $base721 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/' + $base711 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/' + if ($Version.Major -eq 7 -and $Version.Minor -eq 2) { + return @{ + Torch = $base721 + 'torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + Torchvision = $base721 + 'torchvision-0.24.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + Torchaudio = $base721 + 'torchaudio-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + } + } + if ($Version.Major -eq 7 -and $Version.Minor -eq 1) { + return @{ + Torch = $base711 + 'torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + Torchvision = $base711 + 'torchvision-0.24.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + Torchaudio = $base711 + 'torchaudio-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + } + } + return $null +} + # Find Visual Studio Build Tools for cmake -G flag. # Strategy: (1) vswhere, (2) scan filesystem (handles broken vswhere registration). # Returns @{ Generator = "Visual Studio 17 2022"; InstallPath = "C:\..."; Source = "..." } or $null. @@ -534,13 +611,42 @@ if (-not $HasNvidiaSmi) { } } } -if (-not $HasNvidiaSmi) { +# AMD GPU presence via WMI -- always available on Windows, no elevation, +# no dependency on HIP SDK. NVIDIA wins on mixed systems, matching Linux +# install.sh behaviour. +$HasAmdGpu = $false +try { + $videoControllers = Get-CimInstance Win32_VideoController -ErrorAction SilentlyContinue + if ($videoControllers) { + $amdMatches = @($videoControllers | Where-Object { $_.Name -match 'AMD|Radeon' }) + if ($amdMatches.Count -gt 0) { $HasAmdGpu = $true } + } +} catch {} +if ($HasNvidiaSmi) { $HasAmdGpu = $false } + +# Resolve HIP SDK when taking the AMD path. +$HipSdkVersion = $null +if ($HasAmdGpu) { + $HipSdkVersion = Get-HipSdkVersion +} + +if ($HasNvidiaSmi) { + step "gpu" "NVIDIA GPU detected" +} elseif ($HasAmdGpu) { + if ($HipSdkVersion) { + step "gpu" ("AMD GPU detected (HIP SDK {0}.{1})" -f $HipSdkVersion.Major, $HipSdkVersion.Minor) + } else { + Write-Host "" + step "gpu" "AMD GPU detected (HIP SDK missing)" "Yellow" + substep "Install HIP SDK from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" "Yellow" + substep "and re-run this installer." "Yellow" + Write-Host "" + } +} else { Write-Host "" step "gpu" "none (chat-only / GGUF)" "Yellow" - substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" + substep "Training and GPU inference require an NVIDIA or AMD GPU with drivers installed." "Yellow" Write-Host "" -} else { - step "gpu" "NVIDIA GPU detected" } # ============================================ @@ -1359,7 +1465,13 @@ if (Test-Path $VenvDir -PathType Container) { $torchVer = $proc.StandardOutput.ReadToEnd().Trim() $finished = $proc.WaitForExit(30000) if ($finished -and $proc.ExitCode -eq 0 -and $torchVer) { - if ($torchVer -match '\+(cu\d+)') { + # ROCm wheels ship with either "+rocm" (Radeon 7.2.x) or + # "+rocmsdk" (Radeon 7.1.x). A single "rocm" prefix + # match covers both so a ROCm minor update does not trigger + # a pointless venv rebuild. + if ($torchVer -match '\+rocm') { + $installedTorchTag = "rocm" + } elseif ($torchVer -match '\+(cu\d+)') { $installedTorchTag = $Matches[1] } elseif ($torchVer -match '\+cpu') { $installedTorchTag = "cpu" @@ -1380,7 +1492,13 @@ if (Test-Path $VenvDir -PathType Container) { } if (-not $shouldRebuild) { - $expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" } + if ($HasNvidiaSmi) { + $expectedTorchTag = Get-PytorchCudaTag + } elseif ($HasAmdGpu) { + $expectedTorchTag = "rocm" + } else { + $expectedTorchTag = "cpu" + } if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { $shouldRebuild = $true } @@ -1513,11 +1631,56 @@ substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" if ($HasNvidiaSmi) { $CuTag = Get-PytorchCudaTag +} elseif ($HasAmdGpu) { + $CuTag = "rocm" } else { $CuTag = "cpu" } -if ($CuTag -eq "cpu") { +if ($CuTag -eq "rocm") { + if (-not $HipSdkVersion) { + Write-Host "[FAILED] AMD GPU detected but HIP SDK is not installed." -ForegroundColor Red + Write-Host " Download it from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -ForegroundColor Yellow + Write-Host " and re-run this setup." -ForegroundColor Yellow + exit 1 + } + $RocmWheelUrls = Get-RocmWheelUrls -Version $HipSdkVersion + if (-not $RocmWheelUrls) { + Write-Host "[FAILED] AMD HIP SDK $($HipSdkVersion.Major).$($HipSdkVersion.Minor) is not supported." -ForegroundColor Red + Write-Host " Unsloth on Windows requires HIP SDK 7.1 or later. Please update from" -ForegroundColor Yellow + Write-Host " https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -ForegroundColor Yellow + 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. + 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" + } + substep "installing PyTorch ROCm wheels from repo.radeon.com..." + substep "(first wheel is ~780 MB; download may take a few minutes)" + if ($script:UnslothVerbose) { + Fast-Install $RocmWheelUrls.Torch $RocmWheelUrls.Torchvision $RocmWheelUrls.Torchaudio + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install $RocmWheelUrls.Torch $RocmWheelUrls.Torchvision $RocmWheelUrls.Torchaudio | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -ne 0) { + Write-Host "[FAILED] PyTorch ROCm install failed (exit code $torchInstallExit)" -ForegroundColor Red + Write-Host $output -ForegroundColor Red + exit 1 + } + # Triton has no Windows ROCm build; skip the Triton-for-Windows step so + # we do not poison the venv with a package that only targets CUDA. + substep "Triton skipped on Windows AMD (no ROCm build available)" "DarkGray" +} elseif ($CuTag -eq "cpu") { substep "installing PyTorch (CPU-only)..." if ($script:UnslothVerbose) { Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cpu" From c32ff8c6509e44ac347034970cb7d248bac1d331 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Apr 2026 00:02:02 +0000 Subject: [PATCH 02/15] Force pip (skip uv) for bitsandbytes ROCm install Aligns install_python_stack.py with unslothai/unsloth#4966. uv's installer corrupts the bitsandbytes continuous-release_main wheel on ROCm even when the command reports success, leaving the venv with a broken bnb import at runtime. The install.sh fix in #4966 switched to python -m pip install for the Linux shell installer; this commit does the same for the Python installer that runs from unsloth studio update. Changes - pip_install_try and pip_install gain a force_pip: bool = False parameter. When True, the uv attempt is skipped entirely and the call goes straight to python -m pip install via the existing pip_cmd builder. - _ensure_rocm_torch passes force_pip=True for both the bnb pre-release URL install AND the PyPI fallback. Both code paths install bitsandbytes and both are affected by the uv corruption bug, so keeping them consistent matches the gemini-code-assist review comment on #4966 (the fallback in #4966 itself is still uv-backed). - Non-bnb calls (torch install, base packages, extras, etc.) keep the default force_pip=False and continue to prefer uv for speed. --- studio/install_python_stack.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 588671ab11..7a572cf6c9 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -527,6 +527,13 @@ def _ensure_rocm_torch() -> None: # Install bitsandbytes only when torch links against ROCm. Prefers the # continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls # back to PyPI when the pre-release URL is unreachable. + # + # Both installs pass force_pip=True so uv is bypassed entirely. uv's + # installer corrupts the bitsandbytes wheel on ROCm even when the + # command reports success, leaving the venv with a broken bnb import + # at runtime. install.sh has the same fix in unslothai/unsloth#4966. + # We apply it to both the pre-release URL and the PyPI fallback so the + # fix stays consistent regardless of which branch runs. if rocm_torch_ready: # bitsandbytes has no official Windows ROCm wheel # (bitsandbytes-foundation/bitsandbytes#1844), so skip it entirely @@ -544,6 +551,7 @@ def _ensure_rocm_torch() -> None: "--no-deps", _bnb_url, constrain = False, + force_pip = True, ) if not _bnb_installed: print( @@ -560,6 +568,7 @@ def _ensure_rocm_torch() -> None: "--no-deps", _BNB_ROCM_PYPI_FALLBACK, constrain = False, + force_pip = True, ) @@ -859,15 +868,21 @@ def pip_install_try( label: str, *args: str, constrain: bool = True, + force_pip: bool = False, ) -> bool: """Like pip_install but returns False on failure instead of exiting. For optional installs with a follow-up fallback. + + ``force_pip`` skips uv entirely and goes straight to ``python -m pip + install``. Used for wheels that uv installs incorrectly -- notably + the bitsandbytes continuous-release_main wheel on ROCm, see + unslothai/unsloth#4966. """ constraint_args: list[str] = [] if constrain and CONSTRAINTS.is_file(): constraint_args = ["-c", str(CONSTRAINTS)] - if USE_UV: + if USE_UV and not force_pip: cmd = _build_uv_cmd(args) + constraint_args else: cmd = _build_pip_cmd(args) + constraint_args @@ -891,8 +906,15 @@ def pip_install( *args: str, req: Path | None = None, constrain: bool = True, + force_pip: bool = False, ) -> None: - """Build and run a pip install command (uses uv when available, falls back to pip).""" + """Build and run a pip install command (uses uv when available, falls back to pip). + + ``force_pip`` skips the uv attempt entirely and goes straight to + ``python -m pip install``. Use this for wheels that uv installs + incorrectly -- bitsandbytes pre-release ROCm wheels are the known + culprit, see unslothai/unsloth#4966. + """ constraint_args: list[str] = [] if constrain and CONSTRAINTS.is_file(): constraint_args = ["-c", str(CONSTRAINTS)] @@ -910,7 +932,7 @@ def pip_install( req_args = ["-r", str(actual_req)] try: - if USE_UV: + if USE_UV and not force_pip: uv_cmd = _build_uv_cmd(args) + constraint_args + req_args if VERBOSE: print(f" {label}...") @@ -926,7 +948,7 @@ def pip_install( print(result.stdout.decode(errors = "replace")) pip_cmd = _build_pip_cmd(args) + constraint_args + req_args - run(f"{label} (pip)" if USE_UV else label, pip_cmd) + run(f"{label} (pip)" if USE_UV and not force_pip else label, pip_cmd) finally: for temp_req in temp_reqs: temp_req.unlink(missing_ok = True) From f03d6dbe843df5b297da011ae0f20a57f3046bdf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Apr 2026 00:54:10 +0000 Subject: [PATCH 03/15] Correct repo.radeon.com layout comment on _ROCM_WINDOWS_TORCH_WHEELS Earlier comment claimed rocm-rel-6.4.4 uses a "nested layout and alpha version strings". Verified against repo.radeon.com and the reality is more specific: 6.4.4 exposes wheels through a PEP 503 simple index (torch/, torchvision/, torchaudio/ sub-indexes that link back to wheels at the top of the release dir), and the wheels carry alpha plus opaque git-hash build tags like torch-2.8.0a0+gitfc14c65-cp312-cp312-win_amd64.whl which change whenever AMD rebuilds, so they cannot be hardcoded. Also documents that rocm-rel-7.2/ (January) is a distinct release from rocm-rel-7.2.1/ (March) and that the map intentionally routes HIP SDK 7.2.x requests to the newer 7.2.1 wheels because torch bundles its own ROCm runtime. Comment-only change; no behavioural impact. --- studio/install_python_stack.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 7a572cf6c9..8aa050ad53 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -45,10 +45,26 @@ _PYTORCH_WHL_BASE = "https://download.pytorch.org/whl" # Windows AMD ROCm torch wheels live at repo.radeon.com, not download.pytorch.org. # Keyed by HIP SDK (major, minor). Wheels are cp312 only and require the HIP SDK -# to be pre-installed. Older releases (rocm-rel-6.4.4) use a nested layout and -# alpha version strings that are fragile to match on, so only 7.1.x / 7.2.x are -# supported here -- older HIP SDKs fall through with a pointer to the download -# page. +# to be pre-installed. +# +# As of 2026-04, repo.radeon.com/rocm/windows/ contains four release dirs: +# rocm-rel-6.4.4/ -- PEP 503 simple index (torch/, torchvision/, torchaudio/ +# sub-indexes; wheels themselves live at the top of the +# release dir). Wheels carry alpha + git-hash build tags +# like torch-2.8.0a0+gitfc14c65-cp312-cp312-win_amd64.whl, +# so the filename changes whenever AMD rebuilds and we +# cannot hardcode a URL for it. Supporting 6.4.4 would +# require parsing the PEP 503 index at install time -- +# out of scope here; users on that SDK get a "please +# upgrade to 7.1+" error. +# rocm-rel-7.1.1/ -- flat layout, stable `+rocmsdk20251116` date tag. +# rocm-rel-7.2/ -- flat layout, stable `+rocmsdk20260116` date tag. +# rocm-rel-7.2.1/ -- flat layout, stable `+rocm7.2.1` version tag. Newest +# 7.2.x release as of writing; superset of rocm-rel-7.2. +# +# The map below routes HIP SDK 7.2.x -> rocm-rel-7.2.1 wheels (newer, bug +# fixes) rather than rocm-rel-7.2; torch bundles its own ROCm runtime so the +# host SDK point version does not need to match the wheel tag exactly. _ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = { (7, 2): { "torch": ( From 2dce261f4a66afea529d8e8dfa9dfee231d61ffe Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Apr 2026 01:14:23 +0000 Subject: [PATCH 04/15] Install ROCm SDK wheels and relax HIP SDK prerequisite on Windows The previous Windows AMD install was incomplete in two material ways, both surfaced while verifying whether PyTorch has upstream Windows ROCm wheels (it does not -- pytorch.org's get-started page states "ROCm is not available on Windows" and every wheel under download.pytorch.org /whl/rocm7.x is manylinux_2_28_x86_64 only; upstream work is tracked in pytorch/pytorch#159520 and targeted for a future release). repo.radeon.com therefore remains the only source for Windows ROCm torch until that RFC lands, and AMD's install docs at rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html document a two-step pip procedure that we were only implementing half of. Fixing both bugs here so the PR actually produces a working torch.version.hip import on a fresh Windows host. Bug 1: missing ROCm SDK wheels install step AMD's procedure first installs rocm_sdk_core, rocm_sdk_devel, rocm_sdk_libraries_custom, and rocm-.tar.gz (about 1.4 GB total). These wheels ship the ROCm runtime libraries that torch links against at import time. Without them `import torch` fails with missing DLL errors even when the torch wheels themselves are installed. Both 7.1.1 and 7.2.1 need this step; 7.1.1 stamps its SDK wheels with the `0.1.dev0` version string while 7.2.1 uses `7.2.1`. This adds the 4 SDK URLs per release to the wheel map and installs them as Step 1 ahead of the existing torch install (Step 2). Both steps are passed in a single pip call each so pip's dep resolver does not reset torch between wheels (matches AMD's troubleshooting guidance for the same failure mode). Bug 2: HIP SDK was a hard prerequisite, should be optional install.ps1 / setup.ps1 / install_python_stack.py all errored out when $env:HIP_PATH was absent, pointing users at the HIP SDK download page. But the HIP SDK developer toolkit at C:\Program Files\AMD\ROCm\ is for people compiling HIP kernels, not for running PyTorch. AMD's docs list only (a) the AMD graphics driver 26.2.2+ and (b) Python 3.12 as prerequisites. Gating on HIP_PATH was blocking the exact audience that #4280 is about -- regular Radeon users running Unsloth. HIP_PATH is now an optional version hint. When present and valid we use it to select the matching ROCm release; when absent or unsupported (e.g. HIP 6.4) we fall back to the newest stable release (_DEFAULT_WINDOWS_ROCM_VERSION = (7, 2)) and print a visible note pointing at the graphics driver download page. The HIP SDK install prompts have been removed from all three files. Use pip (not uv) for the Radeon wheels Both the SDK and torch install steps call `python -m pip install` directly via the new force_pip=True path added in the previous commit. AMD's documented procedure uses pip, uv has known wheel-corruption issues on similar large ROCm wheels (unslothai/unsloth#4966 for bitsandbytes), and pip is the combination AMD validates. This matches the fix applied to bitsandbytes on Linux ROCm. --- install.ps1 | 133 ++++++++++++++++----- studio/install_python_stack.py | 203 +++++++++++++++++++++++++-------- studio/setup.ps1 | 118 ++++++++++++++----- 3 files changed, 350 insertions(+), 104 deletions(-) diff --git a/install.ps1 b/install.ps1 index 5d70be4b4d..bb52ba58ab 100644 --- a/install.ps1 +++ b/install.ps1 @@ -255,30 +255,53 @@ function Install-UnslothStudio { return $null } - # Map a detected HIP SDK version to Radeon's Windows torch wheels. - # Returns @{ Torch = ...; Torchvision = ...; Torchaudio = ... } or $null - # when the version is unsupported. Wheels are cp312 only. + # Map a ROCm release version to the full Radeon Windows wheel set. + # Returns @{ SdkCore, SdkDevel, SdkLibraries, SdkTarball, Torch, + # Torchvision, Torchaudio } or $null when unsupported. AMD's docs at + # rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html + # require a two-step install: first the rocm_sdk_* wheels (~1.4 GB; + # ship the runtime that torch links against), then torch itself. Both + # are mandatory -- torch import fails with missing DLLs otherwise. + # Wheels are cp312 only. function Get-RocmWheelUrls { param([Parameter(Mandatory = $true)]$Version) $base721 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/' $base711 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/' if ($Version.Major -eq 7 -and $Version.Minor -eq 2) { return @{ - Torch = $base721 + 'torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' - Torchvision = $base721 + 'torchvision-0.24.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' - Torchaudio = $base721 + 'torchaudio-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + SdkCore = $base721 + 'rocm_sdk_core-7.2.1-py3-none-win_amd64.whl' + SdkDevel = $base721 + 'rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl' + SdkLibraries = $base721 + 'rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl' + SdkTarball = $base721 + 'rocm-7.2.1.tar.gz' + Torch = $base721 + 'torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + Torchvision = $base721 + 'torchvision-0.24.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + Torchaudio = $base721 + 'torchaudio-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' } } if ($Version.Major -eq 7 -and $Version.Minor -eq 1) { + # 7.1.1 stamps SDK wheels with `0.1.dev0`; torch gets rocmsdk date tag. return @{ - Torch = $base711 + 'torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' - Torchvision = $base711 + 'torchvision-0.24.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' - Torchaudio = $base711 + 'torchaudio-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + SdkCore = $base711 + 'rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl' + SdkDevel = $base711 + 'rocm_sdk_devel-0.1.dev0-py3-none-win_amd64.whl' + SdkLibraries = $base711 + 'rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl' + SdkTarball = $base711 + 'rocm-0.1.dev0.tar.gz' + Torch = $base711 + 'torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + Torchvision = $base711 + 'torchvision-0.24.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + Torchaudio = $base711 + 'torchaudio-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' } } return $null } + # Default Windows ROCm release when HIP_PATH is absent. HIP_PATH is an + # optional hint, NOT a prerequisite -- regular torch users only need + # an AMD graphics driver (26.2.2+ for 7.2.1) and Python 3.12. PyTorch + # does not publish Windows ROCm wheels on download.pytorch.org (see + # the "ROCm is not available on Windows" note on pytorch.org), so + # repo.radeon.com is the only source until pytorch/pytorch#159520 + # lands upstream Windows ROCm hosting in a future release. + $DefaultWindowsRocmVersion = @{ Major = 7; Minor = 2 } + function New-StudioShortcuts { param( [Parameter(Mandatory = $true)][string]$UnslothExePath @@ -646,9 +669,12 @@ shell.Run cmd, 0, False # users with both cards get the NVIDIA torch path they expect. if ($HasNvidiaSmi) { $HasAmdGpu = $false } - # Resolve HIP SDK when we are taking the AMD path. We still probe here - # even when $HasAmdGpu is $false so the status line is informative on - # mixed NVIDIA+AMD hosts. + # Probe HIP SDK as an OPTIONAL version hint. The HIP SDK developer + # toolkit is NOT a prerequisite for running torch on Windows -- AMD's + # install docs only require the graphics driver (26.2.2+ for 7.2.1) + # and Python 3.12. We use $HipSdkVersion when present to select a + # matching ROCm wheel release; otherwise we fall back to the newest + # stable release ($DefaultWindowsRocmVersion). $HipSdkVersion = $null if ($HasAmdGpu) { $HipSdkVersion = Get-HipSdkVersion @@ -658,11 +684,11 @@ shell.Run cmd, 0, False step "gpu" "NVIDIA GPU detected" } elseif ($HasAmdGpu) { if ($HipSdkVersion) { - step "gpu" ("AMD GPU detected (HIP SDK {0}.{1})" -f $HipSdkVersion.Major, $HipSdkVersion.Minor) + step "gpu" ("AMD GPU detected (HIP SDK {0}.{1} hint)" -f $HipSdkVersion.Major, $HipSdkVersion.Minor) } else { - step "gpu" "AMD GPU detected (HIP SDK missing)" "Yellow" - substep "Install HIP SDK from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" "Yellow" - substep "and re-run this installer." "Yellow" + step "gpu" ("AMD GPU detected (will use rocm-rel-{0}.{1}.x)" -f $DefaultWindowsRocmVersion.Major, $DefaultWindowsRocmVersion.Minor) + substep "HIP SDK not found (optional). Ensure AMD graphics driver is up to date:" "DarkGray" + substep "https://www.amd.com/en/support/download/drivers.html" "DarkGray" } } else { step "gpu" "none (chat-only / GGUF)" "Yellow" @@ -925,12 +951,26 @@ shell.Run cmd, 0, False # The AMD/ROCm path does not use a --index-url; instead it installs # explicit wheel URLs from repo.radeon.com. $TorchIndexUrl stays $null - # on that branch so the NVIDIA/CPU path is visibly bypassed. + # on that branch so the NVIDIA/CPU path is visibly bypassed. HIP_PATH + # is a hint only -- we fall back to $DefaultWindowsRocmVersion when + # it is absent or points at an unsupported version, because the HIP + # SDK is NOT a runtime prerequisite for torch on Windows. $TorchIndexUrl = $null $RocmWheelUrls = $null + $RocmReleaseVersion = $null if ($HasAmdGpu) { if ($HipSdkVersion) { $RocmWheelUrls = Get-RocmWheelUrls -Version $HipSdkVersion + if ($RocmWheelUrls) { + $RocmReleaseVersion = $HipSdkVersion + } + } + if (-not $RocmWheelUrls) { + if ($HipSdkVersion) { + substep ("HIP SDK {0}.{1} is too old; falling back to rocm-rel-{2}.{3}.x" -f $HipSdkVersion.Major, $HipSdkVersion.Minor, $DefaultWindowsRocmVersion.Major, $DefaultWindowsRocmVersion.Minor) "Yellow" + } + $RocmWheelUrls = Get-RocmWheelUrls -Version $DefaultWindowsRocmVersion + $RocmReleaseVersion = $DefaultWindowsRocmVersion } } else { $TorchIndexUrl = Get-TorchIndexUrl @@ -1008,26 +1048,57 @@ shell.Run cmd, 0, False if ($SkipTorch) { substep "skipping PyTorch (--no-torch flag set)." "Yellow" } elseif (-not $RocmWheelUrls) { - if (-not $HipSdkVersion) { - Write-Host "[ERROR] AMD GPU detected but HIP SDK is not installed." -ForegroundColor Red - Write-Host " Download it from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -ForegroundColor Yellow - Write-Host " and re-run this installer." -ForegroundColor Yellow - } else { - Write-Host "[ERROR] AMD HIP SDK $($HipSdkVersion.Major).$($HipSdkVersion.Minor) detected." -ForegroundColor Red - Write-Host " Unsloth requires HIP SDK 7.1 or later on Windows. Please update from" -ForegroundColor Yellow - Write-Host " https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -ForegroundColor Yellow - } + # Should be unreachable because the detection block above + # always falls back to $DefaultWindowsRocmVersion, but guard + # anyway so a future refactor that drops the fallback does + # not install CPU torch on an AMD host. + Write-Host "[ERROR] Could not resolve Windows ROCm wheel URLs." -ForegroundColor Red + Write-Host " This is a bug; please file it at github.com/unslothai/unsloth/issues" -ForegroundColor Yellow return } else { - substep "installing PyTorch ROCm wheels from repo.radeon.com..." - substep "(first wheel is ~780 MB; download may take a few minutes)" + # Verify the venv's Python is 3.12 -- Radeon's wheels are cp312 + # only and pip would fail with a confusing error otherwise. + $venvPyVer = '' + 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") { + 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 + } + + substep ("installing Radeon ROCm wheels for rocm-rel-{0}.{1}.x ..." -f $RocmReleaseVersion.Major, $RocmReleaseVersion.Minor) + substep "Step 1/2: ROCm SDK runtime (~1.4 GB -- this will take a while)" + # Use python -m pip (NOT uv) because (a) AMD's documented + # procedure uses pip, (b) uv has known wheel-corruption issues + # on these big ROCm/bnb wheels (unslothai/unsloth#4966), and + # (c) pip's dep resolver is the combination AMD validates. We + # install all four SDK artefacts in one command so pip does + # not reset torch between them. + $sdkInstallExit = Invoke-InstallCommand { + & $VenvPython -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.SdkCore ` + $RocmWheelUrls.SdkDevel ` + $RocmWheelUrls.SdkLibraries ` + $RocmWheelUrls.SdkTarball + } + if ($sdkInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install ROCm SDK wheels (exit code $sdkInstallExit)" -ForegroundColor Red + Write-Host " Verify your AMD graphics driver is recent and repo.radeon.com is reachable." -ForegroundColor Yellow + return + } + + substep "Step 2/2: PyTorch + torchvision + torchaudio (~820 MB)" $torchInstallExit = Invoke-InstallCommand { - uv pip install --python $VenvPython ` - $RocmWheelUrls.Torch $RocmWheelUrls.Torchvision $RocmWheelUrls.Torchaudio + & $VenvPython -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.Torch ` + $RocmWheelUrls.Torchvision ` + $RocmWheelUrls.Torchaudio } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red - Write-Host " Verify HIP SDK $($HipSdkVersion.Major).$($HipSdkVersion.Minor) is installed and wheels are reachable." -ForegroundColor Yellow + Write-Host " Update your AMD graphics driver: https://www.amd.com/en/support/download/drivers.html" -ForegroundColor Yellow return } } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 8aa050ad53..d087a4a098 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -43,9 +43,23 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { } _PYTORCH_WHL_BASE = "https://download.pytorch.org/whl" -# Windows AMD ROCm torch wheels live at repo.radeon.com, not download.pytorch.org. -# Keyed by HIP SDK (major, minor). Wheels are cp312 only and require the HIP SDK -# to be pre-installed. +# Windows AMD ROCm torch wheels live at repo.radeon.com because PyTorch +# does NOT publish Windows ROCm wheels on download.pytorch.org (and says so +# on pytorch.org/get-started/locally: "ROCm is not available on Windows"). +# Every wheel under download.pytorch.org/whl/rocm{6.4,7.1,7.2}/ is +# manylinux_2_28_x86_64 only. Upstream work to ship Windows ROCm wheels is +# tracked at pytorch/pytorch#159520, targeted for torch 2.10/2.11 but not +# yet delivered. Until then repo.radeon.com is the only source. +# +# AMD's official install procedure +# (rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html) is a +# TWO-STEP pip install: +# Step 1: rocm_sdk_core + rocm_sdk_devel + rocm_sdk_libraries_custom + +# rocm-.tar.gz. These wheels ship the ROCm runtime libraries +# that torch links against at import time. They total about 1.4 +# GB. The HIP SDK developer toolkit (HIP_PATH) is NOT a substitute +# -- torch imports the Python-packaged runtime from rocm_sdk_*. +# Step 2: torch + torchvision + torchaudio. About 820 MB. # # As of 2026-04, repo.radeon.com/rocm/windows/ contains four release dirs: # rocm-rel-6.4.4/ -- PEP 503 simple index (torch/, torchvision/, torchaudio/ @@ -55,18 +69,38 @@ _PYTORCH_WHL_BASE = "https://download.pytorch.org/whl" # so the filename changes whenever AMD rebuilds and we # cannot hardcode a URL for it. Supporting 6.4.4 would # require parsing the PEP 503 index at install time -- -# out of scope here; users on that SDK get a "please -# upgrade to 7.1+" error. -# rocm-rel-7.1.1/ -- flat layout, stable `+rocmsdk20251116` date tag. -# rocm-rel-7.2/ -- flat layout, stable `+rocmsdk20260116` date tag. -# rocm-rel-7.2.1/ -- flat layout, stable `+rocm7.2.1` version tag. Newest -# 7.2.x release as of writing; superset of rocm-rel-7.2. +# out of scope here. +# rocm-rel-7.1.1/ -- flat layout, torch stable `+rocmsdk20251116` tag, +# SDK wheels stamped `0.1.dev0` (pre-release marker). +# rocm-rel-7.2/ -- flat layout, torch stable `+rocmsdk20260116` tag. +# rocm-rel-7.2.1/ -- flat layout, torch stable `+rocm7.2.1` tag, SDK +# wheels stamped `7.2.1`. Newest 7.2.x release as of +# writing; superset of rocm-rel-7.2. # -# The map below routes HIP SDK 7.2.x -> rocm-rel-7.2.1 wheels (newer, bug -# fixes) rather than rocm-rel-7.2; torch bundles its own ROCm runtime so the -# host SDK point version does not need to match the wheel tag exactly. +# HIP SDK detection via HIP_PATH is an OPTIONAL version hint, not a +# prerequisite. Users only need an AMD graphics driver (26.2.2+ for 7.2.1) +# and Python 3.12. When HIP_PATH is absent or points at an unsupported +# version, we default to the newest stable release (7.2.1). _ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = { (7, 2): { + # Step 1: ROCm SDK wheels (ship the runtime torch imports) + "sdk_core": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" + "rocm_sdk_core-7.2.1-py3-none-win_amd64.whl" + ), + "sdk_devel": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" + "rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl" + ), + "sdk_libraries": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" + "rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl" + ), + "sdk_tarball": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" + "rocm-7.2.1.tar.gz" + ), + # Step 2: torch wheels "torch": ( "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" "torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl" @@ -81,6 +115,25 @@ _ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = { ), }, (7, 1): { + # Step 1: ROCm SDK wheels -- note 7.1.1 stamps the SDK wheels with + # `0.1.dev0` while keeping the torch wheels at 2.9.0. + "sdk_core": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/" + "rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl" + ), + "sdk_devel": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/" + "rocm_sdk_devel-0.1.dev0-py3-none-win_amd64.whl" + ), + "sdk_libraries": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/" + "rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl" + ), + "sdk_tarball": ( + "https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/" + "rocm-0.1.dev0.tar.gz" + ), + # Step 2: torch wheels "torch": ( "https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/" "torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl" @@ -95,6 +148,13 @@ _ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = { ), }, } +# Default Windows ROCm release when HIP_PATH is absent or unreadable. Users +# only need a recent AMD graphics driver and Python 3.12 -- the HIP SDK +# 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" ) @@ -373,16 +433,31 @@ def _has_usable_nvidia_gpu() -> bool: def _ensure_rocm_torch_windows() -> None: - """Install Radeon's Windows ROCm torch wheels when an AMD GPU + HIP SDK - are both present. Called from _ensure_rocm_torch(). + """Install Radeon's Windows ROCm SDK + torch wheels when an AMD GPU is + present. Called from _ensure_rocm_torch(). - Silently returns when no AMD GPU is visible, so NVIDIA and CPU-only - Windows hosts are never touched. When an AMD GPU is present but the - HIP SDK is missing or too old, prints a pointer to the HIP SDK - download page and returns without raising -- the Linux helper has the - same shape. NVIDIA takes precedence on mixed AMD+NVIDIA hosts so - install.ps1 and setup.ps1 (which install CUDA torch in that case) - are not clobbered. + AMD's documented install procedure + (rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html) is + a TWO-STEP sequence: first the rocm_sdk_* wheels (which ship the ROCm + runtime libraries torch links against at import time), then the torch + wheels themselves. Both steps are mandatory -- torch import fails + with missing-DLL errors without the SDK wheels, even on a host that + has the HIP SDK developer toolkit installed, because torch imports + the Python-packaged runtime. + + HIP_PATH is treated as an OPTIONAL version hint, not a prerequisite. + Regular users only need the AMD graphics driver (26.2.2+ for 7.2.1) + and Python 3.12, as documented by AMD. When HIP_PATH is missing or + points at an unsupported version, we default to the newest stable + release (_DEFAULT_WINDOWS_ROCM_VERSION) rather than erroring out. + + NVIDIA takes precedence on mixed AMD+NVIDIA hosts. Silently returns + when no AMD GPU is visible so NVIDIA and CPU-only Windows hosts are + never touched. Both pip installs pass force_pip=True because uv's + installer has known problems with these wheels -- matches the fix in + unslothai/unsloth#4966 for bitsandbytes on Linux ROCm, and AMD's own + troubleshooting notes flag pip dep-resolver overwrite scenarios on + this procedure. """ # NVIDIA wins on mixed hosts -- matches the Linux branch and avoids # overwriting a freshly installed CUDA torch with ROCm wheels. @@ -391,27 +466,6 @@ def _ensure_rocm_torch_windows() -> None: if not _has_rocm_gpu_windows(): return - ver = _detect_rocm_version_windows() - if ver is None: - _safe_print( - _red( - " AMD GPU detected but HIP SDK was not found. Install it " - f"from {_HIP_SDK_DOWNLOAD_URL} and re-run setup." - ) - ) - return - - wheels = _ROCM_WINDOWS_TORCH_WHEELS.get(ver) - if wheels is None: - _safe_print( - _red( - f" HIP SDK {ver[0]}.{ver[1]} detected. Unsloth on Windows " - f"requires HIP SDK 7.1 or 7.2. Please update from " - f"{_HIP_SDK_DOWNLOAD_URL}" - ) - ) - return - # 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): @@ -424,20 +478,79 @@ def _ensure_rocm_torch_windows() -> None: ) return + # Prefer HIP_PATH as a version hint when available, but fall back to + # the newest stable release so users without the developer SDK still + # get a working install. + detected = _detect_rocm_version_windows() + ver: tuple[int, int] + using_default = False + if detected is not None and detected in _ROCM_WINDOWS_TORCH_WHEELS: + ver = detected + elif detected is not None: + # Detected but unsupported (e.g. 6.4): fall back to newest with a + # visible notice so the user knows what happened. + _safe_print( + _dim( + f" HIP SDK {detected[0]}.{detected[1]} is too old; " + f"falling back to newest stable wheels " + f"({_DEFAULT_WINDOWS_ROCM_VERSION[0]}.{_DEFAULT_WINDOWS_ROCM_VERSION[1]})" + ) + ) + ver = _DEFAULT_WINDOWS_ROCM_VERSION + using_default = True + else: + ver = _DEFAULT_WINDOWS_ROCM_VERSION + using_default = True + + wheels = _ROCM_WINDOWS_TORCH_WHEELS.get(ver) + if wheels is None: + _safe_print( + _red( + f" No Windows ROCm wheel map for version {ver[0]}.{ver[1]}. " + f"Please file an issue at github.com/unslothai/unsloth/issues." + ) + ) + return + + source = "default" if using_default else "HIP_PATH" _safe_print( _dim( - f" HIP SDK {ver[0]}.{ver[1]} -- installing torch from " - f"repo.radeon.com/rocm/windows/" + f" Installing Radeon ROCm wheels for Windows " + f"(rocm-rel-{ver[0]}.{ver[1]}.x, {source}) from repo.radeon.com" ) ) + if using_default: + _safe_print( + _dim( + f" Ensure your AMD graphics driver is recent; get the " + f"latest from {_AMD_RADEON_DRIVER_URL}" + ) + ) + + # Step 1: ROCm SDK wheels (runtime libraries torch imports). ~1.4 GB + # download on a clean venv. pip_install( - f"ROCm torch (Windows, HIP SDK {ver[0]}.{ver[1]})", + f"ROCm SDK (Windows, {ver[0]}.{ver[1]})", + "--force-reinstall", + "--no-cache-dir", + wheels["sdk_core"], + wheels["sdk_devel"], + wheels["sdk_libraries"], + wheels["sdk_tarball"], + constrain = False, + force_pip = True, + ) + + # Step 2: torch wheels. ~820 MB download. + pip_install( + f"ROCm torch (Windows, {ver[0]}.{ver[1]})", "--force-reinstall", "--no-cache-dir", wheels["torch"], wheels["torchvision"], wheels["torchaudio"], constrain = False, + force_pip = True, ) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index a74b116ff5..f9b9eae11a 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -320,29 +320,45 @@ function Get-HipSdkVersion { return $null } -# Map a detected HIP SDK version to Radeon's Windows torch wheels. -# Returns @{ Torch = ...; Torchvision = ...; Torchaudio = ... } or $null. +# Map a ROCm release version to Radeon's full Windows wheel set. +# Returns @{ SdkCore, SdkDevel, SdkLibraries, SdkTarball, Torch, +# Torchvision, Torchaudio } or $null when unsupported. AMD's install docs +# require a two-step install: first the rocm_sdk_* wheels (~1.4 GB; +# runtime torch links against), then torch itself. Wheels are cp312 only. function Get-RocmWheelUrls { param([Parameter(Mandatory = $true)]$Version) $base721 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/' $base711 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/' if ($Version.Major -eq 7 -and $Version.Minor -eq 2) { return @{ - Torch = $base721 + 'torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' - Torchvision = $base721 + 'torchvision-0.24.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' - Torchaudio = $base721 + 'torchaudio-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + SdkCore = $base721 + 'rocm_sdk_core-7.2.1-py3-none-win_amd64.whl' + SdkDevel = $base721 + 'rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl' + SdkLibraries = $base721 + 'rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl' + SdkTarball = $base721 + 'rocm-7.2.1.tar.gz' + Torch = $base721 + 'torch-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + Torchvision = $base721 + 'torchvision-0.24.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' + Torchaudio = $base721 + 'torchaudio-2.9.1%2Brocm7.2.1-cp312-cp312-win_amd64.whl' } } if ($Version.Major -eq 7 -and $Version.Minor -eq 1) { return @{ - Torch = $base711 + 'torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' - Torchvision = $base711 + 'torchvision-0.24.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' - Torchaudio = $base711 + 'torchaudio-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + SdkCore = $base711 + 'rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl' + SdkDevel = $base711 + 'rocm_sdk_devel-0.1.dev0-py3-none-win_amd64.whl' + SdkLibraries = $base711 + 'rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl' + SdkTarball = $base711 + 'rocm-0.1.dev0.tar.gz' + Torch = $base711 + 'torch-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + Torchvision = $base711 + 'torchvision-0.24.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' + Torchaudio = $base711 + 'torchaudio-2.9.0%2Brocmsdk20251116-cp312-cp312-win_amd64.whl' } } return $null } +# Default Windows ROCm release when HIP_PATH is absent. HIP_PATH is an +# optional hint, not a prerequisite -- torch runtime comes from the +# rocm_sdk wheels, so users only need a graphics driver and Python 3.12. +$DefaultWindowsRocmVersion = @{ Major = 7; Minor = 2 } + # Find Visual Studio Build Tools for cmake -G flag. # Strategy: (1) vswhere, (2) scan filesystem (handles broken vswhere registration). # Returns @{ Generator = "Visual Studio 17 2022"; InstallPath = "C:\..."; Source = "..." } or $null. @@ -624,7 +640,9 @@ try { } catch {} if ($HasNvidiaSmi) { $HasAmdGpu = $false } -# Resolve HIP SDK when taking the AMD path. +# Probe HIP SDK as an optional version hint. Not a prerequisite -- torch +# links against the Python-packaged rocm_sdk wheels, so regular users +# only need the AMD graphics driver (26.2.2+ for 7.2.1) and Python 3.12. $HipSdkVersion = $null if ($HasAmdGpu) { $HipSdkVersion = Get-HipSdkVersion @@ -634,13 +652,11 @@ if ($HasNvidiaSmi) { step "gpu" "NVIDIA GPU detected" } elseif ($HasAmdGpu) { if ($HipSdkVersion) { - step "gpu" ("AMD GPU detected (HIP SDK {0}.{1})" -f $HipSdkVersion.Major, $HipSdkVersion.Minor) + step "gpu" ("AMD GPU detected (HIP SDK {0}.{1} hint)" -f $HipSdkVersion.Major, $HipSdkVersion.Minor) } else { - Write-Host "" - step "gpu" "AMD GPU detected (HIP SDK missing)" "Yellow" - substep "Install HIP SDK from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" "Yellow" - substep "and re-run this installer." "Yellow" - Write-Host "" + step "gpu" ("AMD GPU detected (will use rocm-rel-{0}.{1}.x)" -f $DefaultWindowsRocmVersion.Major, $DefaultWindowsRocmVersion.Minor) + substep "HIP SDK not found (optional). Ensure AMD graphics driver is up to date:" "DarkGray" + substep "https://www.amd.com/en/support/download/drivers.html" "DarkGray" } } else { Write-Host "" @@ -1638,19 +1654,30 @@ if ($HasNvidiaSmi) { } if ($CuTag -eq "rocm") { - if (-not $HipSdkVersion) { - Write-Host "[FAILED] AMD GPU detected but HIP SDK is not installed." -ForegroundColor Red - Write-Host " Download it from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -ForegroundColor Yellow - Write-Host " and re-run this setup." -ForegroundColor Yellow - exit 1 + # HIP_PATH is an optional hint -- fall back to newest stable when + # absent or unsupported, because the HIP developer SDK is NOT a + # runtime prerequisite for torch on Windows (AMD's docs only list + # the graphics driver + Python 3.12). Users without HIP_PATH still + # get a working install. + $RocmWheelUrls = $null + $RocmReleaseVersion = $null + if ($HipSdkVersion) { + $RocmWheelUrls = Get-RocmWheelUrls -Version $HipSdkVersion + if ($RocmWheelUrls) { $RocmReleaseVersion = $HipSdkVersion } } - $RocmWheelUrls = Get-RocmWheelUrls -Version $HipSdkVersion if (-not $RocmWheelUrls) { - Write-Host "[FAILED] AMD HIP SDK $($HipSdkVersion.Major).$($HipSdkVersion.Minor) is not supported." -ForegroundColor Red - Write-Host " Unsloth on Windows requires HIP SDK 7.1 or later. Please update from" -ForegroundColor Yellow - Write-Host " https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html" -ForegroundColor Yellow + if ($HipSdkVersion) { + substep ("HIP SDK {0}.{1} is too old; falling back to rocm-rel-{2}.{3}.x" -f $HipSdkVersion.Major, $HipSdkVersion.Minor, $DefaultWindowsRocmVersion.Major, $DefaultWindowsRocmVersion.Minor) "Yellow" + } + $RocmWheelUrls = Get-RocmWheelUrls -Version $DefaultWindowsRocmVersion + $RocmReleaseVersion = $DefaultWindowsRocmVersion + } + if (-not $RocmWheelUrls) { + Write-Host "[FAILED] Could not resolve Windows ROCm wheel URLs (bug)." -ForegroundColor Red + Write-Host " Please file an issue at github.com/unslothai/unsloth/issues" -ForegroundColor Yellow 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" @@ -1662,14 +1689,49 @@ if ($CuTag -eq "rocm") { 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" } - substep "installing PyTorch ROCm wheels from repo.radeon.com..." - substep "(first wheel is ~780 MB; download may take a few minutes)" + + substep ("installing Radeon ROCm wheels for rocm-rel-{0}.{1}.x from repo.radeon.com..." -f $RocmReleaseVersion.Major, $RocmReleaseVersion.Minor) + substep "Step 1/2: ROCm SDK runtime (~1.4 GB, may take several minutes)" + # Use `python -m pip install` NOT uv for the Radeon wheels. AMD's + # docs specify pip; uv has known issues on similar large ROCm wheels + # (matches the bitsandbytes situation in unslothai/unsloth#4966); and + # using pip directly is the combination AMD validates. if ($script:UnslothVerbose) { - Fast-Install $RocmWheelUrls.Torch $RocmWheelUrls.Torchvision $RocmWheelUrls.Torchaudio + & python -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.SdkCore ` + $RocmWheelUrls.SdkDevel ` + $RocmWheelUrls.SdkLibraries ` + $RocmWheelUrls.SdkTarball + $sdkInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = & python -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.SdkCore ` + $RocmWheelUrls.SdkDevel ` + $RocmWheelUrls.SdkLibraries ` + $RocmWheelUrls.SdkTarball | Out-String + $sdkInstallExit = $LASTEXITCODE + } + if ($sdkInstallExit -ne 0) { + Write-Host "[FAILED] ROCm SDK install failed (exit code $sdkInstallExit)" -ForegroundColor Red + Write-Host $output -ForegroundColor Red + Write-Host " Verify your AMD graphics driver is recent: https://www.amd.com/en/support/download/drivers.html" -ForegroundColor Yellow + exit 1 + } + + substep "Step 2/2: PyTorch + torchvision + torchaudio (~820 MB)" + if ($script:UnslothVerbose) { + & python -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.Torch ` + $RocmWheelUrls.Torchvision ` + $RocmWheelUrls.Torchaudio $torchInstallExit = $LASTEXITCODE $output = "" } else { - $output = Fast-Install $RocmWheelUrls.Torch $RocmWheelUrls.Torchvision $RocmWheelUrls.Torchaudio | Out-String + $output = & python -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.Torch ` + $RocmWheelUrls.Torchvision ` + $RocmWheelUrls.Torchaudio | Out-String $torchInstallExit = $LASTEXITCODE } if ($torchInstallExit -ne 0) { From a128c82b7deb0092e2e6c8bc3ab8541e3152c923 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 11 Apr 2026 01:14:33 +0000 Subject: [PATCH 05/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/install_python_stack.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index d087a4a098..6822b02b3d 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -97,8 +97,7 @@ _ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = { "rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl" ), "sdk_tarball": ( - "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" - "rocm-7.2.1.tar.gz" + "https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/" "rocm-7.2.1.tar.gz" ), # Step 2: torch wheels "torch": ( @@ -152,9 +151,7 @@ _ROCM_WINDOWS_TORCH_WHEELS: dict[tuple[int, int], dict[str, str]] = { # only need a recent AMD graphics driver and Python 3.12 -- the HIP SDK # 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" -) +_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" ) From 7c6086e408b8aaa1b0bd6e366c9491e805769d2d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Apr 2026 01:48:50 +0000 Subject: [PATCH 06/15] Install ROCm SDK + torch in a single pip call on Windows AMD torch's metadata on repo.radeon.com declares `Requires-Dist: rocm[libraries]==`, which cascades to `rocm-sdk-libraries-custom==`. That package does not exist on PyPI at all, and rocm/rocm-sdk-core/rocm-sdk-devel on PyPI are pinned to 0.1.0 (wrong version). AMD's docs show a two-step install but when pip is invoked twice with --force-reinstall on the torch step, the resolver re-evaluates transitive deps, looks for rocm-sdk-libraries-custom on PyPI, fails, and aborts. Fix: pass all 7 Radeon artefacts (4 SDK + 3 torch) in a single pip invocation across install.ps1, studio/setup.ps1, and studio/install_python_stack.py so pip sees the full dep graph upfront and never falls back to PyPI for the rocm-sdk-* chain. Also: - install.ps1 now bootstraps pip/setuptools/wheel into the uv venv before the ROCm install, because rocm-.tar.gz is a Python source distribution pip needs setuptools as its build backend for. - Size estimates are now release-dependent: 7.2.1 is ~2.1 GB, 7.1.1 is ~3.9 GB (its sdk_devel wheel alone is 2.4 GB). - setup.ps1 Fast-Installs pip/setuptools/wheel for the same sdist build-isolation reason. --- install.ps1 | 76 ++++++++++++++++++++++------------ studio/install_python_stack.py | 55 ++++++++++++++++-------- studio/setup.ps1 | 75 ++++++++++++++++++--------------- 3 files changed, 129 insertions(+), 77 deletions(-) diff --git a/install.ps1 b/install.ps1 index bb52ba58ab..a8957520ed 100644 --- a/install.ps1 +++ b/install.ps1 @@ -259,10 +259,11 @@ function Install-UnslothStudio { # Returns @{ SdkCore, SdkDevel, SdkLibraries, SdkTarball, Torch, # Torchvision, Torchaudio } or $null when unsupported. AMD's docs at # rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html - # require a two-step install: first the rocm_sdk_* wheels (~1.4 GB; - # ship the runtime that torch links against), then torch itself. Both - # are mandatory -- torch import fails with missing DLLs otherwise. - # Wheels are cp312 only. + # require a two-step install: first the rocm_sdk_* wheels (~1.3 GB + # for 7.2.1; ~3.2 GB for 7.1.1 whose sdk_devel is 2.4 GB) that ship + # the runtime torch links against, then torch itself (~780 MB for + # 7.2.1, ~692 MB for 7.1.1). Both are mandatory -- torch import + # fails with missing DLLs otherwise. Wheels are cp312 only. function Get-RocmWheelUrls { param([Parameter(Mandatory = $true)]$Version) $base721 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/' @@ -1068,36 +1069,59 @@ shell.Run cmd, 0, False return } - substep ("installing Radeon ROCm wheels for rocm-rel-{0}.{1}.x ..." -f $RocmReleaseVersion.Major, $RocmReleaseVersion.Minor) - substep "Step 1/2: ROCm SDK runtime (~1.4 GB -- this will take a while)" - # Use python -m pip (NOT uv) because (a) AMD's documented - # procedure uses pip, (b) uv has known wheel-corruption issues - # on these big ROCm/bnb wheels (unslothai/unsloth#4966), and - # (c) pip's dep resolver is the combination AMD validates. We - # install all four SDK artefacts in one command so pip does - # not reset torch between them. - $sdkInstallExit = Invoke-InstallCommand { + substep ("installing Radeon ROCm SDK + PyTorch wheels for rocm-rel-{0}.{1}.x" -f $RocmReleaseVersion.Major, $RocmReleaseVersion.Minor) + # Total download size is release-dependent: 7.2.1 is ~2.1 GB + # (1.3 GB SDK + 780 MB torch) and 7.1.1 is ~3.9 GB (3.2 GB + # SDK + 692 MB torch; 7.1.1's sdk_devel alone is 2.4 GB). + # Show the number so a user on a metered connection is not + # surprised. + if ($RocmReleaseVersion.Major -eq 7 -and $RocmReleaseVersion.Minor -eq 1) { + substep "Downloading ~3.9 GB from repo.radeon.com (7.1.1 sdk_devel is 2.4 GB)" + } else { + substep "Downloading ~2.1 GB from repo.radeon.com" + } + # Bootstrap pip + setuptools + wheel into the venv before + # running `& $VenvPython -m pip install`. uv venvs do not + # ship pip by default. setuptools is needed because the + # rocm-.tar.gz artefact is a source distribution that + # pip must build with setuptools as its backend. + substep "bootstrapping pip, setuptools, wheel into venv..." + $bootstrapExit = Invoke-InstallCommand { uv pip install --python $VenvPython pip setuptools wheel } + if ($bootstrapExit -ne 0) { + Write-Host "[ERROR] Failed to bootstrap pip/setuptools/wheel into venv (exit $bootstrapExit)" -ForegroundColor Red + return + } + # Install all 7 Radeon artefacts in a SINGLE pip call. AMD's + # docs document a two-step install, but torch's metadata has + # `Requires-Dist: rocm[libraries]==` which cascades to + # rocm-sdk-libraries-custom== -- a package that does NOT + # exist on PyPI (PyPI has rocm/rocm-sdk-core/rocm-sdk-devel + # only at version 0.1.0, wrong version). Splitting the + # install and using --force-reinstall on the torch step makes + # pip cascade-resolve the rocm dep chain against PyPI and + # fail. Passing every URL in one command gives pip's resolver + # the full dep graph upfront, so --force-reinstall works and + # pip never tries to reach PyPI for the missing packages. + # + # We use `$VenvPython -m pip install` (not uv) because: + # (a) AMD's docs validate pip specifically; + # (b) uv has known wheel-corruption issues on similar large + # ROCm wheels (unslothai/unsloth#4966 for bitsandbytes); + # (c) pip's default build-isolation is required to build the + # rocm-.tar.gz source distribution and uv's pip + # shim does not set this up the same way. + $rocmInstallExit = Invoke-InstallCommand { & $VenvPython -m pip install --no-cache-dir --force-reinstall ` $RocmWheelUrls.SdkCore ` $RocmWheelUrls.SdkDevel ` $RocmWheelUrls.SdkLibraries ` - $RocmWheelUrls.SdkTarball - } - if ($sdkInstallExit -ne 0) { - Write-Host "[ERROR] Failed to install ROCm SDK wheels (exit code $sdkInstallExit)" -ForegroundColor Red - Write-Host " Verify your AMD graphics driver is recent and repo.radeon.com is reachable." -ForegroundColor Yellow - return - } - - substep "Step 2/2: PyTorch + torchvision + torchaudio (~820 MB)" - $torchInstallExit = Invoke-InstallCommand { - & $VenvPython -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.SdkTarball ` $RocmWheelUrls.Torch ` $RocmWheelUrls.Torchvision ` $RocmWheelUrls.Torchaudio } - if ($torchInstallExit -ne 0) { - Write-Host "[ERROR] Failed to install ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red + if ($rocmInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install ROCm wheels (exit code $rocmInstallExit)" -ForegroundColor Red Write-Host " Update your AMD graphics driver: https://www.amd.com/en/support/download/drivers.html" -ForegroundColor Yellow return } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 6822b02b3d..aea1bc085a 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -55,11 +55,23 @@ _PYTORCH_WHL_BASE = "https://download.pytorch.org/whl" # (rocm.docs.amd.com/projects/radeon-ryzen/.../install-pytorch.html) is a # TWO-STEP pip install: # Step 1: rocm_sdk_core + rocm_sdk_devel + rocm_sdk_libraries_custom + -# rocm-.tar.gz. These wheels ship the ROCm runtime libraries -# that torch links against at import time. They total about 1.4 -# GB. The HIP SDK developer toolkit (HIP_PATH) is NOT a substitute -# -- torch imports the Python-packaged runtime from rocm_sdk_*. -# Step 2: torch + torchvision + torchaudio. About 820 MB. +# rocm-.tar.gz. The three wheels ship the ROCm runtime +# libraries torch links against at import time; the tar.gz is a +# tiny (~15 KB) Python sdist for a meta-package named "rocm" +# that requires the other three. The HIP SDK developer toolkit +# (HIP_PATH) is NOT a substitute -- torch imports the +# Python-packaged runtime from rocm_sdk_*. Step 1 download size +# varies dramatically by release: 7.2.1 is about 1.3 GB total +# (sdk_core ~615 MB, sdk_devel ~222 MB, sdk_libraries ~467 MB) +# while 7.1.1 is about 3.2 GB total -- its sdk_devel wheel is a +# massive 2.4 GB (debug symbols / unstripped libraries). +# Step 2: torch + torchvision + torchaudio. About 780 MB for 7.2.1 +# (torch itself is 783 MB; vision/audio are small), 692 MB for +# 7.1.1. +# +# All URLs in _ROCM_WINDOWS_TORCH_WHEELS below were verified live on +# 2026-04-11 with HEAD requests; see temp/windows_amd_tests for the +# verification script. # # As of 2026-04, repo.radeon.com/rocm/windows/ contains four release dirs: # rocm-rel-6.4.4/ -- PEP 503 simple index (torch/, torchvision/, torchaudio/ @@ -524,25 +536,34 @@ def _ensure_rocm_torch_windows() -> None: ) ) - # Step 1: ROCm SDK wheels (runtime libraries torch imports). ~1.4 GB - # download on a clean venv. + # Install all 7 Radeon artefacts (4 SDK + 3 torch) in a SINGLE pip call. + # + # Why one call instead of AMD's documented two steps? torch's metadata + # declares `Requires-Dist: rocm[libraries]==` which cascades to + # `rocm-sdk-libraries-custom==`. That package does NOT exist on + # PyPI; it is only reachable via the direct URL at repo.radeon.com. + # Similarly, `rocm-sdk-core==` on PyPI is stuck at 0.1.0, wrong + # version. If we split the install and use --force-reinstall on the + # torch step, pip's resolver cascades to re-resolve all transitive + # deps, searches PyPI for rocm-sdk-libraries-custom, fails to find + # it, and aborts the whole install. + # + # Passing every URL in one command gives pip's resolver the full dep + # graph upfront. pip picks the right sources, builds the tarball via + # default build isolation, and --force-reinstall works correctly. + # Total download: ~2.1 GB for 7.2.1, ~3.9 GB for 7.1.1 (7.1.1's + # sdk_devel wheel is a massive 2.4 GB -- appears to ship debug + # symbols or unstripped libraries). pip_install( - f"ROCm SDK (Windows, {ver[0]}.{ver[1]})", + f"ROCm SDK + PyTorch (Windows, {ver[0]}.{ver[1]})", "--force-reinstall", "--no-cache-dir", + # SDK artefacts (AMD docs call this Step 1) wheels["sdk_core"], wheels["sdk_devel"], wheels["sdk_libraries"], wheels["sdk_tarball"], - constrain = False, - force_pip = True, - ) - - # Step 2: torch wheels. ~820 MB download. - pip_install( - f"ROCm torch (Windows, {ver[0]}.{ver[1]})", - "--force-reinstall", - "--no-cache-dir", + # torch artefacts (AMD docs call this Step 2) wheels["torch"], wheels["torchvision"], wheels["torchaudio"], diff --git a/studio/setup.ps1 b/studio/setup.ps1 index f9b9eae11a..3500b3caf4 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -323,8 +323,11 @@ function Get-HipSdkVersion { # Map a ROCm release version to Radeon's full Windows wheel set. # Returns @{ SdkCore, SdkDevel, SdkLibraries, SdkTarball, Torch, # Torchvision, Torchaudio } or $null when unsupported. AMD's install docs -# require a two-step install: first the rocm_sdk_* wheels (~1.4 GB; -# runtime torch links against), then torch itself. Wheels are cp312 only. +# require a two-step install: first the rocm_sdk_* wheels (~1.3 GB for +# 7.2.1; ~3.2 GB for 7.1.1) that ship the runtime torch links against, +# then torch itself (~780 MB for 7.2.1, ~692 MB for 7.1.1). Wheels are +# cp312 only. All 14 URLs in the map below were HEAD-verified live on +# 2026-04-11. function Get-RocmWheelUrls { param([Parameter(Mandatory = $true)]$Version) $base721 = 'https://repo.radeon.com/rocm/windows/rocm-rel-7.2.1/' @@ -1690,53 +1693,57 @@ if ($CuTag -eq "rocm") { substep "Re-create the venv with Python 3.12 from https://python.org if pip fails below." "Yellow" } - substep ("installing Radeon ROCm wheels for rocm-rel-{0}.{1}.x from repo.radeon.com..." -f $RocmReleaseVersion.Major, $RocmReleaseVersion.Minor) - substep "Step 1/2: ROCm SDK runtime (~1.4 GB, may take several minutes)" - # Use `python -m pip install` NOT uv for the Radeon wheels. AMD's - # docs specify pip; uv has known issues on similar large ROCm wheels - # (matches the bitsandbytes situation in unslothai/unsloth#4966); and - # using pip directly is the combination AMD validates. - if ($script:UnslothVerbose) { - & python -m pip install --no-cache-dir --force-reinstall ` - $RocmWheelUrls.SdkCore ` - $RocmWheelUrls.SdkDevel ` - $RocmWheelUrls.SdkLibraries ` - $RocmWheelUrls.SdkTarball - $sdkInstallExit = $LASTEXITCODE - $output = "" + substep ("installing Radeon ROCm SDK + PyTorch for rocm-rel-{0}.{1}.x from repo.radeon.com..." -f $RocmReleaseVersion.Major, $RocmReleaseVersion.Minor) + # 7.2.1 total is ~2.1 GB; 7.1.1 is ~3.9 GB (sdk_devel alone is 2.4 GB). + if ($RocmReleaseVersion.Major -eq 7 -and $RocmReleaseVersion.Minor -eq 1) { + substep "Downloading ~3.9 GB (7.1.1 sdk_devel is 2.4 GB -- may take a while)" } else { - $output = & python -m pip install --no-cache-dir --force-reinstall ` + substep "Downloading ~2.1 GB" + } + # setuptools and wheel are required because rocm-.tar.gz is a + # source distribution that pip builds with setuptools.build_meta. + # Ensure both are present in the venv before the install. + if ($script:UnslothVerbose) { + Fast-Install pip setuptools wheel + } else { + Fast-Install pip setuptools wheel | Out-Null + } + # Install all 7 Radeon artefacts in a SINGLE `python -m pip install` + # call. AMD's docs show a two-step install, but torch's metadata + # declares `Requires-Dist: rocm[libraries]==` which cascades to + # rocm-sdk-libraries-custom== -- a package that does NOT exist + # on PyPI. Splitting the install and using --force-reinstall on the + # torch step makes pip cascade-resolve through PyPI and fail. + # Combining into one command gives pip's resolver the full dep graph + # upfront. Use pip (NOT uv) because AMD validates pip, uv has known + # wheel-corruption issues on similar large ROCm wheels (#4966), and + # pip's default build-isolation handles the rocm-.tar.gz sdist. + if ($script:UnslothVerbose) { + & python -m pip install --no-cache-dir --force-reinstall ` $RocmWheelUrls.SdkCore ` $RocmWheelUrls.SdkDevel ` $RocmWheelUrls.SdkLibraries ` - $RocmWheelUrls.SdkTarball | Out-String - $sdkInstallExit = $LASTEXITCODE - } - if ($sdkInstallExit -ne 0) { - Write-Host "[FAILED] ROCm SDK install failed (exit code $sdkInstallExit)" -ForegroundColor Red - Write-Host $output -ForegroundColor Red - Write-Host " Verify your AMD graphics driver is recent: https://www.amd.com/en/support/download/drivers.html" -ForegroundColor Yellow - exit 1 - } - - substep "Step 2/2: PyTorch + torchvision + torchaudio (~820 MB)" - if ($script:UnslothVerbose) { - & python -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.SdkTarball ` $RocmWheelUrls.Torch ` $RocmWheelUrls.Torchvision ` $RocmWheelUrls.Torchaudio - $torchInstallExit = $LASTEXITCODE + $rocmInstallExit = $LASTEXITCODE $output = "" } else { $output = & python -m pip install --no-cache-dir --force-reinstall ` + $RocmWheelUrls.SdkCore ` + $RocmWheelUrls.SdkDevel ` + $RocmWheelUrls.SdkLibraries ` + $RocmWheelUrls.SdkTarball ` $RocmWheelUrls.Torch ` $RocmWheelUrls.Torchvision ` $RocmWheelUrls.Torchaudio | Out-String - $torchInstallExit = $LASTEXITCODE + $rocmInstallExit = $LASTEXITCODE } - if ($torchInstallExit -ne 0) { - Write-Host "[FAILED] PyTorch ROCm install failed (exit code $torchInstallExit)" -ForegroundColor Red + if ($rocmInstallExit -ne 0) { + Write-Host "[FAILED] ROCm SDK + PyTorch install failed (exit code $rocmInstallExit)" -ForegroundColor Red Write-Host $output -ForegroundColor Red + Write-Host " Verify your AMD graphics driver is recent: https://www.amd.com/en/support/download/drivers.html" -ForegroundColor Yellow exit 1 } # Triton has no Windows ROCm build; skip the Triton-for-Windows step so From 7d722faedc3d2df60da982c84e072264caa0e457 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Apr 2026 21:22:45 +0000 Subject: [PATCH 07/15] 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) From b91f47c3b73bbdf9661f7281a987f6a02ca900a1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 12 Apr 2026 21:23:01 +0000 Subject: [PATCH 08/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/install_python_stack.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 52381bfc9a..19d19b1fbe 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -500,8 +500,11 @@ def _ensure_rocm_torch_windows() -> None: # 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 '')"], + [ + sys.executable, + "-c", + "import torch; print(getattr(torch.version,'hip','') or '')", + ], stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, timeout = 30, From d41d593d156b41c8a044d3c8143c1d53b9a2ea48 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Apr 2026 21:45:50 +0000 Subject: [PATCH 09/15] Fix setup.ps1 no-torch crash, idempotency, stale-venv, and Python version guard - setup.ps1: respect UNSLOTH_NO_TORCH when choosing $CuTag so AMD + --no-torch users on Python 3.13 do not crash on the 3.12 version check. - setup.ps1: add torch.version.hip idempotency probe before the ROCm wheel download block. Without this, fresh installs (install.ps1 -> setup.ps1) and every studio update re-downloaded 2.1-3.9 GB of wheels even when ROCm torch was already installed. - setup.ps1: skip venv rebuild for cpu->rocm transitions on AMD hosts. The old behavior deleted the venv then exited with "Run install.ps1 first", breaking the upgrade path for existing CPU-only AMD users. Now keeps the venv and lets the ROCm install block repair torch in-place. - setup.ps1: capture stderr (2>&1) in the non-verbose ROCm pip install so failure diagnostics appear in the error banner. - install_python_stack.py: change the Python 3.12 version check from warn-and-return to sys.exit(1). The old behavior continued to completion and returned exit code 0 without installing ROCm torch. - install_python_stack.py: reorder _ensure_rocm_torch_windows() to put the cheap torch.version.hip probe before the expensive GPU detection subprocess calls, saving ~1-4s per call when ROCm is already installed. - install_python_stack.py: cache _has_rocm_gpu_windows() result so the PowerShell/WMI subprocess is spawned at most once per process. --- studio/install_python_stack.py | 42 ++++++++++++++++++++++------------ studio/setup.ps1 | 36 +++++++++++++++++++++++++---- 2 files changed, 59 insertions(+), 19 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 19d19b1fbe..a764b0f851 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -385,6 +385,9 @@ def _detect_rocm_version_windows() -> tuple[int, int] | None: return None +_HAS_ROCM_GPU_WINDOWS: bool | None = None # module-level cache + + def _has_rocm_gpu_windows() -> bool: """Return True when a Radeon/AMD GPU is visible in WMI Win32_VideoController. @@ -392,8 +395,15 @@ def _has_rocm_gpu_windows() -> bool: HIP SDK -- if we used it to decide whether to prompt the user to install the HIP SDK we would never trigger the prompt on the hosts that need it most. WMI is always available on Windows and needs no elevation. + + Result is cached so repeated calls (steps 2b and 13) do not spawn + a second PowerShell process (~0.5-2 s per call). """ + global _HAS_ROCM_GPU_WINDOWS + if _HAS_ROCM_GPU_WINDOWS is not None: + return _HAS_ROCM_GPU_WINDOWS if not IS_WINDOWS: + _HAS_ROCM_GPU_WINDOWS = False return False ps_cmd = ( "Get-CimInstance Win32_VideoController -ErrorAction SilentlyContinue " @@ -418,7 +428,9 @@ def _has_rocm_gpu_windows() -> bool: continue raw = (result.stdout or "").strip() if raw.isdigit() and int(raw) > 0: + _HAS_ROCM_GPU_WINDOWS = True return True + _HAS_ROCM_GPU_WINDOWS = False return False @@ -486,18 +498,10 @@ def _ensure_rocm_torch_windows() -> None: troubleshooting notes flag pip dep-resolver overwrite scenarios on this procedure. """ - # NVIDIA wins on mixed hosts -- matches the Linux branch and avoids - # overwriting a freshly installed CUDA torch with ROCm wheels. - if _has_usable_nvidia_gpu(): - return - 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. + # Cheap idempotency probe first -- no subprocess spawn needed when + # torch already links against ROCm (common at step 13 and on updates). + # Placed before the expensive GPU-detection calls so the happy-path + # (ROCm already installed) avoids two subprocess spawns entirely. try: _probe = subprocess.run( [ @@ -514,8 +518,16 @@ def _ensure_rocm_torch_windows() -> None: 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. + # NVIDIA wins on mixed hosts -- matches the Linux branch and avoids + # overwriting a freshly installed CUDA torch with ROCm wheels. + if _has_usable_nvidia_gpu(): + return + if not _has_rocm_gpu_windows(): + return + + # Radeon wheels are cp312 only. Hard-exit so the caller (setup.ps1 or + # install.ps1) sees a non-zero exit code instead of continuing with a + # CPU-only torch that silently reports success. if (sys.version_info.major, sys.version_info.minor) != (3, 12): _safe_print( _red( @@ -524,7 +536,7 @@ def _ensure_rocm_torch_windows() -> None: f"Install Python 3.12 from https://python.org and re-run." ) ) - return + sys.exit(1) # Prefer HIP_PATH as a version hint when available, but fall back to # the newest stable release so users without the developer SDK still diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 3461194e61..be8f4081e5 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1519,7 +1519,17 @@ if (Test-Path $VenvDir -PathType Container) { $expectedTorchTag = "cpu" } if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { - $shouldRebuild = $true + # Existing Windows AMD users commonly have a CPU-only venv from + # pre-ROCm installs. Rebuilding from scratch would delete the + # venv and then exit with "Run install.ps1 first" because + # setup.ps1 cannot recreate the venv on its own. Instead, keep + # the venv and let the ROCm install block below repair torch + # in-place -- the end state is the same but no work is lost. + if ($HasAmdGpu -and $installedTorchTag -eq "cpu") { + substep "CPU-only torch detected on AMD host; will repair to ROCm in place..." "Yellow" + } else { + $shouldRebuild = $true + } } } @@ -1648,9 +1658,14 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User') substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" +# --no-torch mode: skip the ROCm wheel path entirely. AMD users running +# install.ps1 --no-torch get a Python 3.13 venv (fine for GGUF); entering +# the ROCm block would hard-fail on the 3.12 version check for no reason. +$_NoTorch = $env:UNSLOTH_NO_TORCH -in @("1", "true", "True", "TRUE") + if ($HasNvidiaSmi) { $CuTag = Get-PytorchCudaTag -} elseif ($HasAmdGpu) { +} elseif ($HasAmdGpu -and -not $_NoTorch) { $CuTag = "rocm" } else { $CuTag = "cpu" @@ -1698,6 +1713,17 @@ if ($CuTag -eq "rocm") { exit 1 } + # Skip the expensive reinstall when ROCm torch is already healthy. + # Mirrors the idempotency guard in _ensure_rocm_torch_windows() -- + # without this, fresh installs (install.ps1 -> setup.ps1) and every + # `unsloth studio update` would re-download 2.1-3.9 GB unnecessarily. + $_existingHip = "" + try { + $_existingHip = (& python -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>$null | Out-String).Trim() + } catch {} + if ($_existingHip) { + substep ("ROCm torch already installed (HIP $_existingHip) -- skipping reinstall") "DarkGray" + } else { substep ("installing Radeon ROCm SDK + PyTorch for rocm-rel-{0}.{1}.x from repo.radeon.com..." -f $RocmReleaseVersion.Major, $RocmReleaseVersion.Minor) # 7.2.1 total is ~2.1 GB; 7.1.1 is ~3.9 GB (sdk_devel alone is 2.4 GB). if ($RocmReleaseVersion.Major -eq 7 -and $RocmReleaseVersion.Minor -eq 1) { @@ -1742,15 +1768,17 @@ if ($CuTag -eq "rocm") { $RocmWheelUrls.SdkTarball ` $RocmWheelUrls.Torch ` $RocmWheelUrls.Torchvision ` - $RocmWheelUrls.Torchaudio | Out-String + $RocmWheelUrls.Torchaudio 2>&1 | Out-String $rocmInstallExit = $LASTEXITCODE } if ($rocmInstallExit -ne 0) { Write-Host "[FAILED] ROCm SDK + PyTorch install failed (exit code $rocmInstallExit)" -ForegroundColor Red Write-Host $output -ForegroundColor Red - Write-Host " Verify your AMD graphics driver is recent: https://www.amd.com/en/support/download/drivers.html" -ForegroundColor Yellow + Write-Host " Possible causes: network error, disk full, or outdated AMD graphics driver." -ForegroundColor Yellow + Write-Host " Update AMD graphics driver: https://www.amd.com/en/support/download/drivers.html" -ForegroundColor Yellow exit 1 } + } # end of: if (-not $_existingHip) # Triton has no Windows ROCm build; skip the Triton-for-Windows step so # we do not poison the venv with a package that only targets CUDA. substep "Triton skipped on Windows AMD (no ROCm build available)" "DarkGray" From f542c231dc7089a5aedc67c10c600cd268981496 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Apr 2026 22:09:27 +0000 Subject: [PATCH 10/15] Fix setup.ps1 version fast-path skipping ROCm repair on up-to-date installs When an existing Windows AMD user runs `unsloth studio update` and the unsloth package is already at the latest version, the version fast-path set $SkipPythonDeps = $true, which bypassed the entire ROCm install block. This left the user permanently stuck on CPU-only torch even though the stale-venv check correctly identified the cpu->rocm mismatch. Fix: track $_NeedRocmRepair flag from the stale-venv check and use it to override $SkipPythonDeps so the ROCm wheel download still runs. Also gate the stale-venv expectedTorchTag on $_NoTorch so --no-torch users on AMD don't see a misleading "will repair to ROCm" message when the actual $CuTag is "cpu". --- studio/setup.ps1 | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index be8f4081e5..0828eb81e4 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1510,10 +1510,15 @@ if (Test-Path $VenvDir -PathType Container) { $shouldRebuild = $true } + # $_NeedRocmRepair is used later to override the version fast-path + # ($SkipPythonDeps) so the ROCm install block still runs even when + # the unsloth package is already at the latest version. + $_NeedRocmRepair = $false + if (-not $shouldRebuild) { if ($HasNvidiaSmi) { $expectedTorchTag = Get-PytorchCudaTag - } elseif ($HasAmdGpu) { + } elseif ($HasAmdGpu -and -not $_NoTorch) { $expectedTorchTag = "rocm" } else { $expectedTorchTag = "cpu" @@ -1525,8 +1530,9 @@ if (Test-Path $VenvDir -PathType Container) { # setup.ps1 cannot recreate the venv on its own. Instead, keep # the venv and let the ROCm install block below repair torch # in-place -- the end state is the same but no work is lost. - if ($HasAmdGpu -and $installedTorchTag -eq "cpu") { + if ($HasAmdGpu -and -not $_NoTorch -and $installedTorchTag -eq "cpu") { substep "CPU-only torch detected on AMD host; will repair to ROCm in place..." "Yellow" + $_NeedRocmRepair = $true } else { $shouldRebuild = $true } @@ -1607,9 +1613,11 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") { $LatestVer = "$($pypiJson.info.version)".Trim() } catch { } - if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) { + if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer) -and -not $_NeedRocmRepair) { step "python" "$_PkgName $InstalledVer is up to date" $SkipPythonDeps = $true + } elseif ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer) -and $_NeedRocmRepair) { + substep "$_PkgName $InstalledVer is current; repairing torch to ROCm..." } elseif ($InstalledVer -and $LatestVer) { substep "$_PkgName $InstalledVer -> $LatestVer available, updating..." } elseif (-not $LatestVer) { From e00eeffe2c42a1b61530dd14c21ccca184f42d2b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Apr 2026 22:27:53 +0000 Subject: [PATCH 11/15] Move $_NoTorch initialization before stale-venv detection in setup.ps1 $_NoTorch was defined at line 1672 but read at lines 1521/1533 in the stale-venv detection block. In PowerShell, reading an undefined variable returns $null, and -not $null evaluates to $true, so the $_NoTorch guards added in the previous commit were silently no-ops. Fix: move the $env:UNSLOTH_NO_TORCH check to before the stale-venv block so --no-torch mode is correctly respected in the torch tag comparison and ROCm repair decision. --- studio/setup.ps1 | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 0828eb81e4..3290f0399d 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1510,6 +1510,11 @@ if (Test-Path $VenvDir -PathType Container) { $shouldRebuild = $true } + # --no-torch: skip ROCm torch repair and wheel install entirely. + # Must be defined here (before stale-venv detection) because the + # expectedTorchTag and $_NeedRocmRepair logic below depend on it. + $_NoTorch = $env:UNSLOTH_NO_TORCH -in @("1", "true", "True", "TRUE") + # $_NeedRocmRepair is used later to override the version fast-path # ($SkipPythonDeps) so the ROCm install block still runs even when # the unsloth package is already at the latest version. @@ -1666,10 +1671,8 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User') substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" -# --no-torch mode: skip the ROCm wheel path entirely. AMD users running -# install.ps1 --no-torch get a Python 3.13 venv (fine for GGUF); entering -# the ROCm block would hard-fail on the 3.12 version check for no reason. -$_NoTorch = $env:UNSLOTH_NO_TORCH -in @("1", "true", "True", "TRUE") +# $_NoTorch was already initialized earlier (before stale-venv detection) +# so it is available here for the $CuTag selection. if ($HasNvidiaSmi) { $CuTag = Get-PytorchCudaTag From ed41a254398bffbf721d0d9dd0b0c2bd33d52e76 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Apr 2026 22:52:35 +0000 Subject: [PATCH 12/15] Trim verbose review-fix comments Shorten multi-line comments added during review iterations to 1-2 lines each. Remove redundant explanations where the code is self-evident. --- install.ps1 | 4 +--- studio/install_python_stack.py | 17 ++++++----------- studio/setup.ps1 | 31 +++++++++---------------------- 3 files changed, 16 insertions(+), 36 deletions(-) diff --git a/install.ps1 b/install.ps1 index c1dbb2ddbd..cebe7069c0 100644 --- a/install.ps1 +++ b/install.ps1 @@ -784,9 +784,7 @@ 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 + 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. + # AMD + torch requires Python 3.12 (Radeon wheels are cp312 only). if ($HasAmdGpu -and -not $SkipTorch) { $PythonPreferred = @("3.12") $PythonVersion = "3.12" diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a764b0f851..d84bc5a064 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -438,10 +438,8 @@ 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. + # nvidia-smi.exe is often not on PATH on Windows; match the + # fallback paths in install.ps1 / setup.ps1. _candidates = [ os.path.join( os.environ.get("ProgramFiles", r"C:\Program Files"), @@ -498,10 +496,8 @@ def _ensure_rocm_torch_windows() -> None: troubleshooting notes flag pip dep-resolver overwrite scenarios on this procedure. """ - # Cheap idempotency probe first -- no subprocess spawn needed when - # torch already links against ROCm (common at step 13 and on updates). - # Placed before the expensive GPU-detection calls so the happy-path - # (ROCm already installed) avoids two subprocess spawns entirely. + # Cheap idempotency probe first -- skip GPU detection when torch + # already links against ROCm (common at step 13 and on updates). try: _probe = subprocess.run( [ @@ -525,9 +521,8 @@ def _ensure_rocm_torch_windows() -> None: if not _has_rocm_gpu_windows(): return - # Radeon wheels are cp312 only. Hard-exit so the caller (setup.ps1 or - # install.ps1) sees a non-zero exit code instead of continuing with a - # CPU-only torch that silently reports success. + # Radeon wheels are cp312 only. Hard-exit so the caller sees a + # non-zero code instead of silently leaving CPU-only torch. if (sys.version_info.major, sys.version_info.minor) != (3, 12): _safe_print( _red( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 3290f0399d..8be29cb727 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1510,14 +1510,9 @@ if (Test-Path $VenvDir -PathType Container) { $shouldRebuild = $true } - # --no-torch: skip ROCm torch repair and wheel install entirely. - # Must be defined here (before stale-venv detection) because the - # expectedTorchTag and $_NeedRocmRepair logic below depend on it. + # Must be defined before stale-venv detection (expectedTorchTag depends on it). $_NoTorch = $env:UNSLOTH_NO_TORCH -in @("1", "true", "True", "TRUE") - - # $_NeedRocmRepair is used later to override the version fast-path - # ($SkipPythonDeps) so the ROCm install block still runs even when - # the unsloth package is already at the latest version. + # Overrides $SkipPythonDeps so the ROCm block runs even when unsloth is current. $_NeedRocmRepair = $false if (-not $shouldRebuild) { @@ -1529,12 +1524,9 @@ if (Test-Path $VenvDir -PathType Container) { $expectedTorchTag = "cpu" } if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { - # Existing Windows AMD users commonly have a CPU-only venv from - # pre-ROCm installs. Rebuilding from scratch would delete the - # venv and then exit with "Run install.ps1 first" because - # setup.ps1 cannot recreate the venv on its own. Instead, keep - # the venv and let the ROCm install block below repair torch - # in-place -- the end state is the same but no work is lost. + # Keep the venv and let the ROCm block below repair torch + # in-place; rebuilding would delete it and exit since + # setup.ps1 cannot recreate a venv on its own. if ($HasAmdGpu -and -not $_NoTorch -and $installedTorchTag -eq "cpu") { substep "CPU-only torch detected on AMD host; will repair to ROCm in place..." "Yellow" $_NeedRocmRepair = $true @@ -1671,8 +1663,7 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User') substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" -# $_NoTorch was already initialized earlier (before stale-venv detection) -# so it is available here for the $CuTag selection. +# $_NoTorch initialized earlier (before stale-venv detection). if ($HasNvidiaSmi) { $CuTag = Get-PytorchCudaTag @@ -1707,9 +1698,8 @@ if ($CuTag -eq "rocm") { exit 1 } - # 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. + # Radeon wheels are cp312 only. Hard-stop instead of downloading + # multi-GB wheels that pip will reject. try { $venvPyVer = (& python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')" 2>$null | Out-String).Trim() } catch { $venvPyVer = "" } @@ -1724,10 +1714,7 @@ if ($CuTag -eq "rocm") { exit 1 } - # Skip the expensive reinstall when ROCm torch is already healthy. - # Mirrors the idempotency guard in _ensure_rocm_torch_windows() -- - # without this, fresh installs (install.ps1 -> setup.ps1) and every - # `unsloth studio update` would re-download 2.1-3.9 GB unnecessarily. + # Skip reinstall when ROCm torch is already healthy (idempotency guard). $_existingHip = "" try { $_existingHip = (& python -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>$null | Out-String).Trim() From 267601e52a4b222c3a72d8a25aead9c7dc5b59f3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 13:44:23 +0000 Subject: [PATCH 13/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/install_python_stack.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 1d4eeef599..e252955cde 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1180,11 +1180,14 @@ def _has_usable_nvidia_gpu() -> bool: for _candidate in ( os.path.join( os.environ.get("ProgramFiles", r"C:\Program Files"), - "NVIDIA Corporation", "NVSMI", "nvidia-smi.exe", + "NVIDIA Corporation", + "NVSMI", + "nvidia-smi.exe", ), os.path.join( os.environ.get("SystemRoot", r"C:\Windows"), - "System32", "nvidia-smi.exe", + "System32", + "nvidia-smi.exe", ), ): if os.path.isfile(_candidate): From 10bc8aa538198d52924dd2f5fd9d84368ecb5f45 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 27 Jul 2026 23:17:58 +0000 Subject: [PATCH 14/15] Keep probing for nvidia-smi after an unusable one on PATH _has_usable_nvidia_gpu gated its Windows fixed-location fallback on the PATH lookup missing, not on the GPU check failing. A stale or driverless nvidia-smi exits non-zero listing nothing, so the search stopped there and the function reported no NVIDIA GPU even with a working driver binary under NVSMI or System32. That answer routes a mixed AMD iGPU plus NVIDIA dGPU Windows host into _ensure_rocm_torch() and replaces its CUDA stack with ROCm wheels. install.ps1 and studio/setup.ps1 already do the right thing: both call Test-NvidiaSmiHasGpu on the PATH result and fall through to the two fixed paths when it fails, with the same reasoning recorded at install.ps1:1708 ("a stale/driverless nvidia-smi can exit 0 while listing no GPU"). This brings the Python helper to the same rule: collect the candidates, then take the first that lists a GPU. Reproduced with real stub executables through the real subprocess call, before and after: PATH exe fixed-location exe before after absent working True True stale working False True working - True True none none False False Only the stale row changes. An AMD-only host with a leftover nvidia-smi still gets False, so it is not denied the ROCm wheels. tests/studio/install/test_nvidia_smi_candidate_probing.py pins all four rows plus the CUDA_VISIBLE_DEVICES cases: 1 failed / 8 passed before, 9 passed after. --- studio/install_python_stack.py | 52 +++++--- .../test_nvidia_smi_candidate_probing.py | 119 ++++++++++++++++++ 2 files changed, 151 insertions(+), 20 deletions(-) create mode 100644 tests/studio/install/test_nvidia_smi_candidate_probing.py diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e252955cde..d27fff0d82 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1175,9 +1175,32 @@ def _has_usable_nvidia_gpu() -> bool: cvd = os.environ.get("CUDA_VISIBLE_DEVICES") if cvd is not None and cvd.strip() in ("", "-1"): return False - exe = shutil.which("nvidia-smi") - if not exe and IS_WINDOWS: - for _candidate in ( + def _lists_a_gpu(exe: str) -> bool: + 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 + + # Try every candidate until one lists a GPU, rather than committing to the + # first executable found. A stale or driverless nvidia-smi on PATH exits + # non-zero listing nothing; stopping there would report the host as + # NVIDIA-free and route it into _ensure_rocm_torch() even though a working + # driver binary sits at a fixed location. install.ps1 and setup.ps1 both + # gate their fallback on the GPU check failing, not on the PATH lookup + # missing, so mirror that. + candidates = [] + _path_exe = shutil.which("nvidia-smi") + if _path_exe: + candidates.append(_path_exe) + if IS_WINDOWS: + candidates.extend(( os.path.join( os.environ.get("ProgramFiles", r"C:\Program Files"), "NVIDIA Corporation", @@ -1189,23 +1212,12 @@ def _has_usable_nvidia_gpu() -> bool: "System32", "nvidia-smi.exe", ), - ): - if os.path.isfile(_candidate): - exe = _candidate - break - if exe: - try: - result = subprocess.run( - [exe, "-L"], - stdout = subprocess.PIPE, - stderr = subprocess.DEVNULL, - text = True, - timeout = 10, - ) - if result.returncode == 0 and "GPU " in result.stdout: - return True - except Exception: - pass + )) + for _candidate in candidates: + if _candidate != _path_exe and not os.path.isfile(_candidate): + continue + if _lists_a_gpu(_candidate): + return True # Fallback: the NVIDIA driver exposes one subdirectory per GPU under # /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state. if sys.platform != "win32": diff --git a/tests/studio/install/test_nvidia_smi_candidate_probing.py b/tests/studio/install/test_nvidia_smi_candidate_probing.py new file mode 100644 index 0000000000..c33c551417 --- /dev/null +++ b/tests/studio/install/test_nvidia_smi_candidate_probing.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""_has_usable_nvidia_gpu must keep probing after an unusable nvidia-smi. + +A stale or driverless nvidia-smi on PATH exits non-zero and lists no GPU. +Treating the first executable found as the answer reports a mixed AMD+NVIDIA +Windows host as NVIDIA-free, which routes it into _ensure_rocm_torch() and +replaces a working CUDA stack with ROCm wheels. install.ps1 and setup.ps1 both +gate their fixed-location fallback on the GPU check failing rather than on the +PATH lookup missing; this pins the Python helper to the same rule. + +The stubs are real executables run through the real subprocess call, so the +test exercises the actual control flow rather than a mocked return value. +""" + +import importlib.util +import os +import pathlib +import sys +import types + +import pytest + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +_STUDIO = _REPO_ROOT / "studio" +_SRC = _STUDIO / "install_python_stack.py" + +_STALE = 'echo "No devices were found"; exit 9' +_WORKING = 'echo "GPU 0: NVIDIA H100 (UUID: GPU-abc)"; exit 0' + + +def _load_module(): + # install_python_stack imports backend.utils.wheel_utils, which resolves + # only with studio/ on sys.path. That is how the installer invokes it. + if str(_STUDIO) not in sys.path: + sys.path.insert(0, str(_STUDIO)) + spec = importlib.util.spec_from_file_location("_ips_probe_under_test", _SRC) + module = importlib.util.module_from_spec(spec) + sys.modules["_ips_probe_under_test"] = module + spec.loader.exec_module(module) + return module + + +def _write_stub(path: pathlib.Path, body: str) -> None: + path.parent.mkdir(parents = True, exist_ok = True) + # Not /usr/bin/env: PATH is narrowed to the stub directory below, so env + # would not find an interpreter. + path.write_text("#!/bin/bash\n" + body + "\n") + path.chmod(0o755) + + +@pytest.fixture +def probe(tmp_path, monkeypatch): + """Run _has_usable_nvidia_gpu as if on Windows, with stubbed nvidia-smi.""" + + def _run( + path_smi: str | None, + fixed_smi: str | None, + cuda_visible_devices: str | None = None, + ) -> bool: + path_dir = tmp_path / "pathbin" + path_dir.mkdir(exist_ok = True) + if path_smi is not None: + _write_stub(path_dir / "nvidia-smi", path_smi) + program_files = tmp_path / "ProgramFiles" + if fixed_smi is not None: + _write_stub( + program_files / "NVIDIA Corporation" / "NVSMI" / "nvidia-smi.exe", + fixed_smi, + ) + monkeypatch.setenv("PATH", str(path_dir)) + monkeypatch.setenv("ProgramFiles", str(program_files)) + monkeypatch.setenv("SystemRoot", str(tmp_path / "Windows")) + if cuda_visible_devices is None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + else: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda_visible_devices) + module = _load_module() + monkeypatch.setattr(module, "IS_WINDOWS", True) + # A real NVIDIA host has /proc/driver/nvidia/gpus, and the helper's + # Linux-only fallback would then answer True for every case and mask + # what the Windows path did. Present as win32 to isolate it. + monkeypatch.setattr(module, "sys", types.SimpleNamespace(platform = "win32")) + return module._has_usable_nvidia_gpu() + + return _run + + +def test_stale_path_nvidia_smi_still_reaches_the_fixed_locations(probe): + # The regression: a driverless nvidia-smi on PATH used to end the search. + assert probe(_STALE, _WORKING) is True + + +def test_absent_path_nvidia_smi_reaches_the_fixed_locations(probe): + assert probe(None, _WORKING) is True + + +def test_working_path_nvidia_smi_is_enough(probe): + assert probe(_WORKING, None) is True + + +def test_no_nvidia_smi_anywhere_reports_no_gpu(probe): + assert probe(None, None) is False + + +def test_stale_everywhere_reports_no_gpu(probe): + # Every candidate answering "no GPU" must stay False, or an AMD-only host + # with a leftover nvidia-smi would be denied the ROCm wheels. + assert probe(_STALE, _STALE) is False + + +@pytest.mark.parametrize("hidden", ["", "-1", " "]) +def test_cuda_visible_devices_hidden_wins_over_a_working_probe(probe, hidden): + assert probe(_WORKING, _WORKING, cuda_visible_devices = hidden) is False + + +def test_cuda_visible_devices_listing_a_device_does_not_block_detection(probe): + assert probe(_WORKING, None, cuda_visible_devices = "0") is True From 534ba3bfa25755eff85b3fbaa43d3ee86d507f89 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 23:19:34 +0000 Subject: [PATCH 15/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/install_python_stack.py | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index d27fff0d82..65530a3f00 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1175,6 +1175,7 @@ def _has_usable_nvidia_gpu() -> bool: cvd = os.environ.get("CUDA_VISIBLE_DEVICES") if cvd is not None and cvd.strip() in ("", "-1"): return False + def _lists_a_gpu(exe: str) -> bool: try: result = subprocess.run( @@ -1200,19 +1201,21 @@ def _has_usable_nvidia_gpu() -> bool: if _path_exe: candidates.append(_path_exe) if IS_WINDOWS: - candidates.extend(( - os.path.join( - os.environ.get("ProgramFiles", r"C:\Program Files"), - "NVIDIA Corporation", - "NVSMI", - "nvidia-smi.exe", - ), - os.path.join( - os.environ.get("SystemRoot", r"C:\Windows"), - "System32", - "nvidia-smi.exe", - ), - )) + candidates.extend( + ( + os.path.join( + os.environ.get("ProgramFiles", r"C:\Program Files"), + "NVIDIA Corporation", + "NVSMI", + "nvidia-smi.exe", + ), + os.path.join( + os.environ.get("SystemRoot", r"C:\Windows"), + "System32", + "nvidia-smi.exe", + ), + ) + ) for _candidate in candidates: if _candidate != _path_exe and not os.path.isfile(_candidate): continue