ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Windows ROCm install: idempotency guard, nvidia-smi fallback, Python version gates

Browse source 
- 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.
This commit is contained in:
Daniel Han 2026-04-12 21:22:45 +00:00
commit 7d722faedc
3 changed files with 62 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

22
install.ps1
View file

@ -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

44
studio/install_python_stack.py
View file

@ -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):

19
studio/setup.ps1
View file

@ -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)