feat: add Windows AMD ROCm PyTorch wheel installation
install_python_stack.py: - Add _ROCM_WINDOWS_WHEEL_BASE and _ROCM_WINDOWS_RELEASES constants pointing to AMD repo.radeon.com (ROCm 7.2 -> torch 2.9.1+rocm7.2.1) - Extend _ensure_rocm_torch() with a Windows branch: detects ROCm via _has_rocm_gpu() / _detect_rocm_version(), requires Python 3.12 (cp312 is the only ABI AMD publishes for Windows), installs the direct wheel URL from repo.radeon.com install.ps1: - Capture ROCmVersion during AMD detection via hipconfig --version / amd-smi version (needed for wheel URL selection) - After Get-TorchIndexUrl, add an AMD wheel override block: when HasROCm and Python 3.12 detected, set ROCmTorchWheelUrl to AMD wheel URL - Expand torch install branch to handle ROCmTorchWheelUrl with uv pip install --force-reinstall --no-cache-dir
This commit is contained in:
LeoBorcherding
parent
5d5ae56566
commit
270b2dd9b0
2 changed files with 134 additions and 14 deletions
65
install.ps1
65
install.ps1
|
|
@ -1209,6 +1209,7 @@ shell.Run cmd, 0, False
|
|||
# ── AMD ROCm detection (Windows) — mirrors setup.ps1 ──
|
||||
$HasROCm = $false
|
||||
$ROCmGpuLabel = $null
|
||||
$ROCmVersion = $null
|
||||
if (-not $HasNvidiaSmi) {
|
||||
$hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue
|
||||
if ($hipinfoExe) {
|
||||
|
|
@ -1244,6 +1245,29 @@ shell.Run cmd, 0, False
|
|||
if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name }
|
||||
} catch {}
|
||||
}
|
||||
# Capture ROCm version for wheel selection (hipconfig, then amd-smi)
|
||||
if ($HasROCm) {
|
||||
$hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue
|
||||
if ($hipConfigExe) {
|
||||
try {
|
||||
$hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String
|
||||
if ($LASTEXITCODE -eq 0 -and $hipVerOut -match '(\d+\.\d+)') {
|
||||
$ROCmVersion = $Matches[1]
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
if (-not $ROCmVersion) {
|
||||
$amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue
|
||||
if ($amdSmiVer) {
|
||||
try {
|
||||
$smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String
|
||||
if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') {
|
||||
$ROCmVersion = $Matches[1]
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if ($HasNvidiaSmi) {
|
||||
|
|
@ -1281,12 +1305,39 @@ shell.Run cmd, 0, False
|
|||
return "$baseUrl/cu126"
|
||||
}
|
||||
$TorchIndexUrl = Get-TorchIndexUrl
|
||||
$TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl
|
||||
|
||||
# ── AMD Windows ROCm wheel override ──
|
||||
# AMD publishes direct torch wheels for Windows (cp312 only) at repo.radeon.com.
|
||||
# When the HIP SDK is present and Python 3.12 is in use, swap in the AMD wheel
|
||||
# URL and clear $TorchIndexUrl so the standard --index-url path is skipped.
|
||||
$ROCmTorchWheelUrl = $null
|
||||
if ($HasROCm -and -not $SkipTorch) {
|
||||
$pyMajMin = if ($DetectedPython) { ($DetectedPython.Version -split '\.')[0..1] -join '.' } else { "" }
|
||||
if ($pyMajMin -eq "3.12") {
|
||||
$amdWheelBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.radeon.com/rocm/windows" }
|
||||
if ($ROCmVersion -and $ROCmVersion -match '^7\.2') {
|
||||
$ROCmTorchWheelUrl = "$amdWheelBase/rocm-rel-7.2.1/torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl"
|
||||
$TorchIndexUrl = $null
|
||||
}
|
||||
if ($ROCmTorchWheelUrl) {
|
||||
substep "AMD ROCm $ROCmVersion (Python 3.12) -- AMD Windows torch wheel selected" "Cyan"
|
||||
} elseif ($ROCmVersion) {
|
||||
substep "No AMD Windows torch wheel for ROCm $ROCmVersion -- falling back to CPU-only PyTorch" "Yellow"
|
||||
} else {
|
||||
substep "ROCm version unknown -- falling back to CPU-only PyTorch" "Yellow"
|
||||
}
|
||||
} else {
|
||||
substep "AMD Windows ROCm wheels require Python 3.12 (detected: $pyMajMin) -- using CPU-only PyTorch" "Yellow"
|
||||
substep "To enable ROCm training, reinstall with Python 3.12." "Yellow"
|
||||
}
|
||||
}
|
||||
|
||||
$TorchIndexFamily = Get-TauriTorchIndexFamily $(if ($ROCmTorchWheelUrl) { "rocm7.2" } else { $TorchIndexUrl })
|
||||
$GpuBranch = Get-TauriGpuBranch $TorchIndexFamily
|
||||
Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version
|
||||
|
||||
# ── Print CPU-only hint when no GPU detected ──
|
||||
if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") {
|
||||
if (-not $SkipTorch -and -not $ROCmTorchWheelUrl -and $TorchIndexUrl -like "*/cpu") {
|
||||
Write-Host ""
|
||||
if ($HasROCm -or $ROCmGpuLabel) {
|
||||
substep "Installing CPU-only PyTorch (ROCm wheels require the HIP SDK)." "Yellow"
|
||||
|
|
@ -1364,9 +1415,17 @@ shell.Run cmd, 0, False
|
|||
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
|
||||
}
|
||||
}
|
||||
} elseif ($TorchIndexUrl) {
|
||||
} elseif ($TorchIndexUrl -or $ROCmTorchWheelUrl) {
|
||||
if ($SkipTorch) {
|
||||
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
|
||||
} elseif ($ROCmTorchWheelUrl) {
|
||||
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
|
||||
substep "installing PyTorch (AMD ROCm $ROCmVersion)..."
|
||||
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-cache-dir $ROCmTorchWheelUrl }
|
||||
if ($torchInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit)
|
||||
}
|
||||
} else {
|
||||
Write-TauriLog "STEP" "Installing PyTorch"
|
||||
substep "installing PyTorch ($TorchIndexUrl)..."
|
||||
|
|
|
|||
|
|
@ -57,6 +57,17 @@ _PYTORCH_WHL_BASE = (
|
|||
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
|
||||
).rstrip("/")
|
||||
|
||||
# AMD Windows ROCm wheels — repo.radeon.com (cp312 only; AMD does not publish
|
||||
# Windows ROCm wheels for other Python versions)
|
||||
_ROCM_WINDOWS_WHEEL_BASE = (
|
||||
os.environ.get("UNSLOTH_ROCM_WINDOWS_MIRROR")
|
||||
or "https://repo.radeon.com/rocm/windows"
|
||||
).rstrip("/")
|
||||
# Maps (major, minor) → (release_folder, torch_version_string)
|
||||
_ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, str]] = {
|
||||
(7, 2): ("rocm-rel-7.2.1", "2.9.1+rocm7.2.1"),
|
||||
}
|
||||
|
||||
# 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.
|
||||
|
|
@ -241,20 +252,70 @@ def _has_usable_nvidia_gpu() -> bool:
|
|||
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).
|
||||
On Linux x86_64: uses pytorch.org ROCm wheel index tags.
|
||||
On Windows (cp312 only): uses AMD's repo.radeon.com direct wheel releases.
|
||||
No-op on macOS, non-x86_64 Linux, NVIDIA-primary hosts, or when torch
|
||||
already links against HIP.
|
||||
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:
|
||||
if IS_MACOS:
|
||||
return
|
||||
|
||||
if IS_WINDOWS:
|
||||
# AMD only publishes Windows ROCm wheels for Python 3.12 (cp312)
|
||||
if sys.version_info[:2] != (3, 12):
|
||||
print(
|
||||
f" ROCm torch on Windows requires Python 3.12 "
|
||||
f"(current: {sys.version_info[0]}.{sys.version_info[1]}) -- skipping"
|
||||
)
|
||||
return
|
||||
if _has_usable_nvidia_gpu():
|
||||
return
|
||||
if not _has_rocm_gpu():
|
||||
return
|
||||
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 ROCm torch
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
ver = _detect_rocm_version()
|
||||
if ver is None:
|
||||
print(" ROCm detected but version unreadable -- skipping torch reinstall")
|
||||
return
|
||||
entry = next(
|
||||
((rt, tv) for (maj, mn), (rt, tv) in sorted(_ROCM_WINDOWS_RELEASES.items(), reverse = True)
|
||||
if ver >= (maj, mn)),
|
||||
None,
|
||||
)
|
||||
if entry is None:
|
||||
print(
|
||||
f" No AMD Windows torch wheel for ROCm {ver[0]}.{ver[1]} -- skipping"
|
||||
)
|
||||
return
|
||||
rel_tag, torch_ver = entry
|
||||
wheel_url = (
|
||||
f"{_ROCM_WINDOWS_WHEEL_BASE}/{rel_tag}/"
|
||||
f"torch-{torch_ver}-cp312-cp312-win_amd64.whl"
|
||||
)
|
||||
print(f" ROCm {ver[0]}.{ver[1]} (Windows) -- installing torch from {wheel_url}")
|
||||
pip_install(
|
||||
f"ROCm torch (Windows, {rel_tag})",
|
||||
"--force-reinstall",
|
||||
"--no-cache-dir",
|
||||
wheel_url,
|
||||
constrain = False,
|
||||
)
|
||||
return
|
||||
|
||||
# ── Linux x86_64 path ──────────────────────────────────────────────────────
|
||||
# PyTorch only publishes ROCm wheels for linux_x86_64; skip aarch64 / arm64
|
||||
# to avoid a missing-wheel error on `unsloth studio update`.
|
||||
if platform.machine().lower() not in {"x86_64", "amd64"}:
|
||||
return
|
||||
# NVIDIA takes precedence on mixed hosts -- but only if an actual GPU is usable
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue