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.
This commit is contained in:
Daniel Han
parent
5a017ecfea
commit
d2d5f90175
3 changed files with 66 additions and 10 deletions
34
install.ps1
34
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) {
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue