diff --git a/install.ps1 b/install.ps1 index 7e4e8b4aac..edec79b29d 100644 --- a/install.ps1 +++ b/install.ps1 @@ -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) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index ceb73951d8..0aa6116a08 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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 diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 4b22579c3b..ffbef0bac8 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -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) {