From d41d593d156b41c8a044d3c8143c1d53b9a2ea48 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Apr 2026 21:45:50 +0000 Subject: [PATCH] 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"