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"