From d2d5f90175960c00e2e2759c92ce166004ac127f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 09:16:19 +0000 Subject: [PATCH] install: complete pinned-index handling for ROCm/Windows edge cases Follow-ups to the override work flagged in review: - install.ps1: a pinned gfx*/rocm>=7.2 index previously skipped the AMD reroute that sets the torch>=2.11 floor, so the generic install used torch>=2.4,<2.11 and could resolve the known-bad _grouped_mm wheel. Route a pinned ROCm index through the ROCm install path with the 2.11 floor + companions, and guard the companion-spec lookup so a skipped reroute block cannot null-deref. - studio/setup.ps1: the stale-venv check compared the installed flavor (cuXXX/cpu, with +rocm misread as cpu) against the raw pinned leaf (gfx1151 / rocm6.4), so a correct pinned ROCm venv was always marked stale. Classify +rocm wheels as the generic 'rocm' flavor and normalize a pinned rocm*/gfx* leaf to 'rocm' before comparing (cu* stays specific so cu126-vs-cu128 still rebuilds). - install_python_stack.py: _ensure_cuda_torch now also reinstalls from a pinned CUDA index when the venv carries a CPU wheel (headless CPU-venv-to-CUDA cross-install via 'studio update'), not only when it finds a ROCm build. - tests: parity assertions already cover all four installers honoring the override. --- install.ps1 | 34 ++++++++++++++++++++++++++++++---- studio/install_python_stack.py | 25 ++++++++++++++++++++----- studio/setup.ps1 | 17 ++++++++++++++++- 3 files changed, 66 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 2a929c81d5..7e4e8b4aac 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2043,6 +2043,8 @@ exit 0 # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. $ROCmIndexUrl = $null $ROCmTorchFloor = $null + $PinnedRocmVisionSpec = $null + $PinnedRocmAudioSpec = $null if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ @@ -2093,6 +2095,30 @@ exit 0 } } + # An explicit gfx*/rocm pin skips the auto-reroute above, but the generic + # CPU/CUDA install below would use torch>=2.4,<2.11 and pull a known-bad wheel + # on the gfx115x/gfx120x/rocm>=7.2 indexes (the torch._C._grouped_mm null-ptr + # bug). Route a pinned ROCm index through the ROCm install path with the same + # 2.11 floor/companions the unpinned reroute derives from the gfx arch. + if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) { + $_pinLeaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLower() + $_pinRocm211 = $false + if ($_pinLeaf -match '^rocm(\d+)\.(\d+)') { + $_pinRocm211 = ([int]$Matches[1] -gt 7) -or ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -ge 2) + } + if ($_pinLeaf -like 'gfx*' -or $_pinRocm211) { + $ROCmIndexUrl = $TorchIndexUrl + $ROCmTorchFloor = "torch>=2.11.0,<2.12.0" + $PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0" + $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0" + substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan" + } elseif ($_pinLeaf -like 'rocm*') { + # Older rocm (<=7.1) ships torch <2.11; route via the ROCm path with the + # default floor so the pinned family resolves its own wheels. + $ROCmIndexUrl = $TorchIndexUrl + } + } + if ($ROCmIndexUrl) { $TorchIndexFamily = "rocm" } else { @@ -2205,8 +2231,8 @@ exit 0 $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin the companions to match $torchSpec; bare names can resolve an # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { # Transient AMD-index failure: fall back to a CPU base so the install @@ -2326,8 +2352,8 @@ exit 0 $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin companions like the fresh ROCm path (bare names can pull an # ABI-incompatible torchvision/torchaudio from the per-arch index). - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 37c3c7d449..ceb73951d8 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1073,16 +1073,31 @@ def _ensure_cuda_torch() -> None: _marker_lines = [ line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() ] - if not _marker_lines or _marker_lines[-1] != "hip": + if not _marker_lines: + return + _marker = _marker_lines[-1] + # Reinstall CUDA torch when the venv carries a ROCm build on an NVIDIA host + # (the poisoning signature), OR when an explicit CUDA index is pinned but the + # venv still has a CPU wheel. The latter is the headless CPU-venv-to-CUDA + # cross-install (`studio update` with UNSLOTH_TORCH_INDEX_FAMILY=cu128): the + # update path preserves torch rather than preinstalling it from install.sh, so + # without this the explicit CUDA pin stays ineffective. A healthy CUDA torch, + # or a CPU wheel with no CUDA pin, is deliberate and left alone. + _pin = _explicit_torch_index_url() + _pinned_cuda = bool(_pin) and _pin.rstrip("/").rsplit("/", 1)[-1].lower().startswith("cu") + if _marker == "hip": + _why = "torch is a ROCm build on an NVIDIA host" + elif _marker == "cpu" and _pinned_cuda: + _why = "torch is a CPU build but an explicit CUDA index is pinned" + else: return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is index_url = _detect_cuda_torch_index_url() _torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC print( - f" torch is a ROCm build on an NVIDIA host -- reinstalling " - f"CUDA torch from {index_url}\n" - f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch " - f"on a mixed AMD+NVIDIA host)" + f" {_why} -- reinstalling CUDA torch from {index_url}\n" + f" (set UNSLOTH_TORCH_BACKEND=rocm or cpu to keep a deliberate " + f"non-CUDA torch)" ) pip_install( "CUDA torch repair", diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 907367a69a..4b22579c3b 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2540,6 +2540,11 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode if ($finished -and $proc.ExitCode -eq 0 -and $torchVer) { if ($torchVer -match '\+(cu\d+)') { $installedTorchTag = $Matches[1] + } elseif ($torchVer -match '\+rocm') { + # Any +rocm / gfx wheel -> generic "rocm" flavor. The exact ROCm + # version is repaired later by install_python_stack.py; here we + # only need the flavor so a correct ROCm venv is not marked stale. + $installedTorchTag = "rocm" } elseif ($torchVer -match '\+cpu') { $installedTorchTag = "cpu" } else { @@ -2560,7 +2565,17 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode if (-not $shouldRebuild) { $_pinnedIdx = Get-PinnedTorchIndexUrl - $expectedTorchTag = if ($_pinnedIdx) { Get-TorchIndexLeaf $_pinnedIdx } elseif ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" } + if ($_pinnedIdx) { + $_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx + # Normalize a pinned rocm*/gfx* leaf to the generic "rocm" flavor so it + # compares against the installed +rocm wheel (also "rocm"); cu*/cpu + # leaves stay specific so a cu126-vs-cu128 mismatch still rebuilds. + $expectedTorchTag = if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { "rocm" } else { $_pinLeaf } + } elseif ($HasNvidiaSmi) { + $expectedTorchTag = Get-PytorchCudaTag + } else { + $expectedTorchTag = "cpu" + } if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { $shouldRebuild = $true }