install: finish pinned ROCm/CUDA edge cases on Windows + repair path
Follow-ups to the previous round: - studio/setup.ps1: a pinned gfx*/rocm>=7.2 index now routes through the ROCm install path with the 2.11 floor + companions (it previously fell through to the CUDA branch with bare torch/torchvision/torchaudio against the ROCm index). The CPU/CUDA fallback index is forced to the CPU wheel index when a ROCm index is active, so a failed pinned-ROCm install does not retry the ROCm mirror. - studio/setup.ps1: the stale-venv check no longer treats an unrecognized pinned URL leaf (e.g. a PEP 503 mirror ending in /simple) as a torch flavor tag, which was marking a correct venv stale; cu*/cpu/rocm/gfx leaves are still compared. - install.ps1: the post-failure CPU fallback uses an explicit CPU index instead of , which for a pinned ROCm index was the ROCm mirror itself (so the 'fallback' just retried the failing index and aborted the installer). - install_python_stack.py: _ensure_cuda_torch now also reinstalls when the venv's CUDA family differs from a pinned one (installed cu126 vs pinned cu128), not only CPU->CUDA; the probe reports the installed cuXXX tag for the comparison.
This commit is contained in:
Daniel Han
parent
d2d5f90175
commit
4a5baba8ce
3 changed files with 63 additions and 15 deletions
|
|
@ -2242,7 +2242,11 @@ exit 0
|
|||
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
|
||||
# torch>= range, so without it uv would keep the ROCm build and only swap
|
||||
# the companions -- a mismatched venv the flavor-repair block won't fix.
|
||||
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
|
||||
# Use an explicit CPU index: in the unpinned AMD path $TorchIndexUrl is
|
||||
# already */cpu, but for a pinned ROCm index it IS the ROCm mirror, so
|
||||
# reusing it here would just retry the failing ROCm index, not fall back.
|
||||
$CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" }
|
||||
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $CpuFallbackIndexUrl }
|
||||
if ($torchInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
|
||||
|
|
|
|||
|
|
@ -1053,11 +1053,13 @@ def _ensure_cuda_torch() -> None:
|
|||
sys.executable,
|
||||
"-c",
|
||||
(
|
||||
"import torch; "
|
||||
"import torch, re; "
|
||||
"hip = getattr(torch.version, 'hip', '') or ''; "
|
||||
"cuda = getattr(torch.version, 'cuda', '') or ''; "
|
||||
"ver = getattr(torch, '__version__', '').lower(); "
|
||||
"print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))"
|
||||
"m = re.search(r'\\+(cu\\d+)', ver); "
|
||||
"marker = 'hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'); "
|
||||
"print(marker + '|' + (m.group(1) if m else ''))"
|
||||
),
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
|
|
@ -1075,22 +1077,26 @@ def _ensure_cuda_torch() -> None:
|
|||
]
|
||||
if not _marker_lines:
|
||||
return
|
||||
_marker = _marker_lines[-1]
|
||||
_marker, _, _installed_cu = _marker_lines[-1].partition("|")
|
||||
# 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.
|
||||
# (the poisoning signature), or when an explicit CUDA index is pinned but the
|
||||
# venv has the wrong family -- a CPU wheel, or a different cuXXX than pinned.
|
||||
# This covers the headless 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 an explicit CUDA pin
|
||||
# stays ineffective. A healthy CUDA torch matching the pin, 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")
|
||||
_pin_leaf = _pin.rstrip("/").rsplit("/", 1)[-1].lower() if _pin else ""
|
||||
_pinned_cuda = _pin_leaf.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"
|
||||
elif _marker == "cuda" and _pinned_cuda and _installed_cu and _installed_cu != _pin_leaf:
|
||||
_why = f"torch is {_installed_cu} but the pinned CUDA index is {_pin_leaf}"
|
||||
else:
|
||||
return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is
|
||||
return # healthy CUDA torch matching the pin, or a deliberate CPU wheel
|
||||
|
||||
index_url = _detect_cuda_torch_index_url()
|
||||
_torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
|
||||
|
|
|
|||
|
|
@ -2565,18 +2565,29 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
|
|||
|
||||
if (-not $shouldRebuild) {
|
||||
$_pinnedIdx = Get-PinnedTorchIndexUrl
|
||||
$_expectedKnown = $true
|
||||
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 }
|
||||
if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') {
|
||||
$expectedTorchTag = "rocm"
|
||||
} elseif ($_pinLeaf -like 'cu*' -or $_pinLeaf -eq 'cpu') {
|
||||
$expectedTorchTag = $_pinLeaf
|
||||
} else {
|
||||
# Custom index whose final segment is not a torch flavor (e.g. a
|
||||
# PEP 503 mirror ending in /simple). We cannot infer the flavor, so
|
||||
# trust the pinned URL and do not rebuild on a bogus tag comparison.
|
||||
$_expectedKnown = $false
|
||||
$expectedTorchTag = $installedTorchTag
|
||||
}
|
||||
} elseif ($HasNvidiaSmi) {
|
||||
$expectedTorchTag = Get-PytorchCudaTag
|
||||
} else {
|
||||
$expectedTorchTag = "cpu"
|
||||
}
|
||||
if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
|
||||
if ($_expectedKnown -and $installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
|
||||
$shouldRebuild = $true
|
||||
}
|
||||
}
|
||||
|
|
@ -2828,12 +2839,39 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu
|
|||
}
|
||||
}
|
||||
|
||||
# A pinned gfx*/rocm index skips the auto-reroute above; route it through the
|
||||
# ROCm install path with the same floor/companions the unpinned AMD path uses
|
||||
# (mirrors install.ps1). Otherwise the CUDA branch below installs bare torch /
|
||||
# torchvision / torchaudio from the ROCm index and resolves a known-bad <2.11
|
||||
# wheel or ABI-mismatched companions for gfx115x / gfx120x / rocm>=7.2.
|
||||
if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) {
|
||||
$_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl
|
||||
$_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 = $PinnedTorchIndexUrl
|
||||
$ROCmTorchSpec = "torch>=2.11.0,<2.12.0"
|
||||
$ROCmVisionSpec = "torchvision>=0.26.0,<0.27.0"
|
||||
$ROCmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
|
||||
substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchSpec" "Cyan"
|
||||
} elseif ($_pinLeaf -like 'rocm*') {
|
||||
$ROCmIndexUrl = $PinnedTorchIndexUrl
|
||||
$ROCmTorchSpec = "torch"
|
||||
$ROCmVisionSpec = "torchvision"
|
||||
$ROCmAudioSpec = "torchaudio"
|
||||
}
|
||||
}
|
||||
|
||||
$PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
|
||||
|
||||
# A full UNSLOTH_TORCH_INDEX_URL pin is used verbatim; a family pin already set
|
||||
# $CuTag, so $PyTorchWhlBase/$CuTag is the requested family index. The CPU/CUDA
|
||||
# install branches below pull from this instead of re-joining mirror + tag.
|
||||
$TorchInstallIndexUrl = if ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
|
||||
# A pinned ROCm install goes through $ROCmIndexUrl; if that fails, the fallback
|
||||
# must use the CPU wheel index, not retry the ROCm mirror left in the pin.
|
||||
$TorchInstallIndexUrl = if ($ROCmIndexUrl) { "$PyTorchWhlBase/cpu" } elseif ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
|
||||
|
||||
$ROCmCpuFallback = $false
|
||||
if ($ROCmIndexUrl) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue