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.
This commit is contained in:
Daniel Han
parent
b91f47c3b7
commit
d41d593d15
2 changed files with 59 additions and 19 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue