install: address #6692 review round 5 (ROCm/CPU pin edge cases)

setup.ps1: - Stale-venv check: treat an AMD/ROCm host (HasROCm or a resolved gfx arch) with no explicit pin as expecting "rocm", not "cpu", so a healthy +rocm venv is not flagged stale (which made installer-managed setup exit and direct update rebuild). - Pinned-ROCm install failure now routes into the force-reinstall CPU branch: CuTag stays the rocm/gfx leaf on failure, so the condition also checks ROCmCpuFallback; otherwise the CUDA branch installed from the CPU index without --force-reinstall and kept the partial ROCm torch. - Explicit ROCm pin compare no longer collapses gfx*/rocm* to a generic "rocm": it compares the +rocmX.Y version (and the torch 2.11 line for gfx pins) so changing the pinned family (e.g. rocm6.4 -> gfx1151) rebuilds and applies it. install_python_stack.py: - _ensure_rocm_torch: an explicit ROCm wheel-index pin now bypasses the NVIDIA-present / no-AMD-GPU / unreadable-ROCm gates (headless/container/CI cross-install), mirroring the explicit-CUDA-pin bypass in _ensure_cuda_torch. - Add _ensure_cpu_torch: an explicit CPU pin (FAMILY=cpu or /cpu URL) now has a repair path that reinstalls CPU torch over an existing CUDA/ROCm build on a standalone update (which skips install.sh's flavor enforcement). install.sh: - Pin torchvision/torchaudio companions alongside torch for the rocm7.2 / per-gfx index and the Strix reroute (those AMD indexes publish companions independently and a bare name can resolve a torch-2.12-built wheel, an ABI mismatch).
2026-06-27 09:11:23 +00:00 · 2026-06-27 09:11:23 +00:00 · c290cac455
commit c290cac455
parent f6c5e46dad
3 changed files with 179 additions and 27 deletions
--- a/install.sh
+++ b/install.sh
@ -1862,6 +1862,14 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t
        TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
    fi
 fi
+# Companion (torchvision/torchaudio) constraints. Bare by default: the pytorch.org
+# cu*/cpu/rocmX.Y indexes are curated so uv resolves an ABI-consistent trio from a
+# bare name. They are pinned alongside TORCH_CONSTRAINT only for the torch-2.11
+# AMD paths (rocm7.2 / per-gfx index / Strix), where AMD publishes each wheel
+# independently and can ship a newer torchvision/torchaudio (built against torch
+# 2.12) before removing the 2.11-matched one -- see the rocm7.2/gfx case below.
+TORCHVISION_CONSTRAINT="torchvision"
+TORCHAUDIO_CONSTRAINT="torchaudio"

 # ── Resolve repo root (for --local installs) ──
 _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
@ -2435,9 +2443,16 @@ esac
 # 2.11.0 -- adjust the constraint to allow it. This also covers a pinned full-URL
 # or family override (e.g. UNSLOTH_TORCH_INDEX_URL=.../gfx1151) that returns early
 # above and so never hits the Strix reroute that otherwise raises this constraint.
+# Pin the companions to the matching 2.11 range too: the per-gfx index publishes
+# torchvision/torchaudio independently and a bare name can resolve a 2.12-built
+# wheel (ABI mismatch). Matches setup.ps1's *FloorMap and _ROCM_TORCH_PKG_SPECS.
 # All other ROCm tags and CUDA stay within <2.11.0.
 case "$TORCH_INDEX_URL" in
-    */rocm7.2|*/gfx*) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
+    */rocm7.2|*/gfx*)
+        TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
+        TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
+        TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
+        ;;
 esac

 # Auto-detect GPU for AMD ROCm based
@ -2524,6 +2539,10 @@ case "$TORCH_INDEX_URL" in
            done
            TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/"
            TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
+            # Pin companions to the 2.11 range (per-gfx index publishes them
+            # independently); mirrors the rocm7.2/gfx case above.
+            TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
+            TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
            _amd_gpu_radeon=false
        fi
        ;;
@ -2712,7 +2731,7 @@ if [ "$_MIGRATED" = true ]; then
                if [ -z "$_has_hip" ]; then
                    substep "repairing ROCm torch (overwritten by dependency resolution)..."
                    run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
-                        "$TORCH_CONSTRAINT" torchvision torchaudio \
+                        "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
                        --index-url "$TORCH_INDEX_URL" \
                        --force-reinstall
                fi
@ -2838,7 +2857,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
                   [ "$_radeon_versions_match" != true ]; then
                    substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
                    run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
-                        "$TORCH_CONSTRAINT" torchvision torchaudio \
+                        "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
                        --index-url "$TORCH_INDEX_URL"
                else
                    substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
@ -2861,18 +2880,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
            else
                substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
                run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
-                    "$TORCH_CONSTRAINT" torchvision torchaudio \
+                    "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
                    --index-url "$TORCH_INDEX_URL"
            fi
        else
            substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
            run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
-                "$TORCH_CONSTRAINT" torchvision torchaudio \
+                "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
                --index-url "$TORCH_INDEX_URL"
        fi
    else
        substep "installing PyTorch ($TORCH_INDEX_URL)..."
-        run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
+        run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
            --index-url "$TORCH_INDEX_URL"
    fi
    # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
@ -2932,7 +2951,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
                if [ -z "$_has_hip" ]; then
                    substep "repairing ROCm torch (overwritten by dependency resolution)..."
                    run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
-                        "$TORCH_CONSTRAINT" torchvision torchaudio \
+                        "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
                        --index-url "$TORCH_INDEX_URL" \
                        --force-reinstall
                fi
@ -2974,7 +2993,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
           && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
            substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
            run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
-                "$TORCH_CONSTRAINT" torchvision torchaudio \
+                "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
                --index-url "$TORCH_INDEX_URL" \
                --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
            _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
@ -2986,7 +3005,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
            substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
            substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
            substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
-            substep "[WARN]   uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
+            substep "[WARN]   uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
        fi
    fi
 fi
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@ -1005,6 +1005,19 @@ def _explicit_rocm_torch_index_url() -> "str | None":
    return url if leaf.startswith(("rocm", "gfx")) else None


+def _explicit_cpu_torch_index_url() -> "str | None":
+    """The pinned wheel index URL when it names the CPU family (leaf == cpu), else None.
+
+    An explicit CPU pin (UNSLOTH_TORCH_INDEX_FAMILY=cpu or a URL ending in /cpu)
+    is authoritative -- see _ensure_cpu_torch.
+    """
+    url = _explicit_torch_index_url()
+    if url is None:
+        return None
+    leaf = url.rstrip("/").rsplit("/", 1)[-1].lower()
+    return url if leaf == "cpu" else None
+
+
 def _ensure_cuda_torch() -> None:
    """Repair a venv whose torch is a ROCm build on an NVIDIA host.

@ -1118,6 +1131,70 @@ def _ensure_cuda_torch() -> None:
    )


+def _ensure_cpu_torch() -> None:
+    """Reinstall CPU torch when an explicit CPU pin is set but the venv has a GPU build.
+
+    Counterpart to _ensure_cuda_torch / _ensure_rocm_torch for the explicit-CPU
+    case (UNSLOTH_TORCH_INDEX_FAMILY=cpu or a URL ending in /cpu). Those helpers
+    both treat a CPU backend as a skip signal, so on a standalone `unsloth studio
+    update` -- which does not run install.sh's post-install flavor enforcement --
+    an existing CUDA/ROCm torch satisfies the version constraint and is never
+    replaced, ignoring the authoritative CPU pin. Only fires for an EXPLICIT pin:
+    a CPU backend that came from auto-detection (genuine CPU host via install.sh)
+    already installed CPU wheels, so there is nothing to repair.
+    """
+    if NO_TORCH:
+        return
+    pin = _explicit_cpu_torch_index_url()
+    if pin is None:
+        return
+
+    # Classify the installed torch family. A non-zero exit means torch is missing
+    # or un-importable; the base install step handles that, so leave it alone.
+    try:
+        probe = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                (
+                    "import torch, re; "
+                    "hip = getattr(torch.version, 'hip', '') or ''; "
+                    "cuda = getattr(torch.version, 'cuda', '') or ''; "
+                    "ver = getattr(torch, '__version__', '').lower(); "
+                    "gpu = bool(hip) or 'rocm' in ver or bool(cuda) or bool(re.search(r'\\+cu\\d+', ver)); "
+                    "print('gpu' if gpu else 'cpu')"
+                ),
+            ],
+            stdout = subprocess.PIPE,
+            stderr = subprocess.DEVNULL,
+            timeout = 90,
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return
+    if probe.returncode != 0:
+        return
+    _lines = [line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip()]
+    if not _lines or _lines[-1] != "gpu":
+        return  # already CPU (or unreadable) -- nothing to repair
+
+    print(
+        f"   torch is a GPU build but an explicit CPU index is pinned -- "
+        f"reinstalling CPU torch from {pin}"
+    )
+    # The pytorch.org /cpu index is curated, so a bare trio resolves consistently.
+    pip_install(
+        "CPU torch repair",
+        "--force-reinstall",
+        "--no-cache-dir",
+        "torch",
+        "torchvision",
+        "torchaudio",
+        "--index-url",
+        pin,
+        constrain = False,
+    )
+
+
 def _ensure_rocm_torch() -> None:
    """Reinstall torch with ROCm wheels when the venv received CPU-only torch.

@ -1248,22 +1325,37 @@ def _ensure_rocm_torch() -> None:
    # ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ──
    if platform.machine().lower() not in {"x86_64", "amd64"}:
        return
-    # NVIDIA takes precedence on mixed hosts -- but only if a GPU is usable
-    if _has_usable_nvidia_gpu():
-        return
-    # Use _has_rocm_gpu() (rocminfo / amd-smi GPU data rows) as the
-    # authoritative "is this an AMD ROCm host?" signal. The old gate required
-    # /opt/rocm or hipcc to exist, which breaks runtime-only ROCm installs
-    # (minimal package-managed installs, Radeon software) that ship
-    # amd-smi/rocminfo without /opt/rocm or hipcc, leaving `unsloth studio
-    # update` unable to repair a CPU-only venv on those systems.
-    if not _has_rocm_gpu():
-        return  # no AMD GPU visible
+    # An explicit ROCm wheel-index pin (UNSLOTH_TORCH_INDEX_URL/_FAMILY naming a
+    # rocm*/gfx* leaf) commits to ROCm wheels regardless of which GPU is visible
+    # here -- the headless / container / CI cross-install case. Mirror
+    # _ensure_cuda_torch's explicit-pin bypass: skip the NVIDIA-present / no-AMD /
+    # unreadable-ROCm gates so the pinned index is honoured. Without this, a
+    # standalone `studio update` with an explicit ROCm pin on an NVIDIA-only or
+    # GPU-less box returned here and left the CPU/CUDA torch in place.
+    _rocm_pin = _explicit_rocm_torch_index_url()
+    if _rocm_pin is None:
+        # NVIDIA takes precedence on mixed hosts -- but only if a GPU is usable.
+        if _has_usable_nvidia_gpu():
+            return
+        # Use _has_rocm_gpu() (rocminfo / amd-smi GPU data rows) as the
+        # authoritative "is this an AMD ROCm host?" signal. The old gate required
+        # /opt/rocm or hipcc to exist, which breaks runtime-only ROCm installs
+        # (minimal package-managed installs, Radeon software) that ship
+        # amd-smi/rocminfo without /opt/rocm or hipcc, leaving `unsloth studio
+        # update` unable to repair a CPU-only venv on those systems.
+        if not _has_rocm_gpu():
+            return  # no AMD GPU visible

    ver = _detect_rocm_version()
    if ver is None:
-        print("   ROCm detected but version unreadable -- skipping torch reinstall")
-        return
+        if _rocm_pin is None:
+            print("   ROCm detected but version unreadable -- skipping torch reinstall")
+            return
+        # Explicit pin: the pinned index leaf (not the host ROCm version) drives
+        # the install below, so a missing/unreadable host ROCm version is fine.
+        # Use a sentinel so the version-gated Strix reroute (skipped for explicit
+        # pins anyway) and any ver comparisons stay well-defined.
+        ver = (0, 0)

    # Probe whether torch already links against HIP (ROCm already working).
    # Do NOT skip for CUDA-only builds: they are unusable on AMD-only hosts
@ -2237,6 +2329,7 @@ def install_python_stack() -> int:
        _progress(_torch_step_label("check"))
        _ensure_cuda_torch()
        _ensure_rocm_torch()
+        _ensure_cpu_torch()

    # Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
    # version or unknown ROCm version).
@ -2427,6 +2520,7 @@ def install_python_stack() -> int:
        _progress(_torch_step_label("final"))
        _ensure_cuda_torch()
        _ensure_rocm_torch()
+        _ensure_cpu_torch()

    # 14. Final check (silent; third-party conflicts are expected)
    subprocess.run(
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@ -2568,11 +2568,39 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
        $_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.
+            # cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds.
            if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') {
-                $expectedTorchTag = "rocm"
+                # Do NOT collapse a pinned ROCm/gfx leaf to a generic "rocm": that
+                # would match any installed +rocm wheel and mask a pin change from
+                # one ROCm family to another (e.g. rocm6.4 -> gfx1151, or rocm6.4
+                # -> rocm6.3), leaving the requested index unapplied. Compare what
+                # the wheel tag exposes -- the ROCm version (+rocmX.Y) and whether
+                # the index serves the torch 2.11 line (gfx* / rocm>=7.2 do; older
+                # rocm does not). gfx pins carry no rocm version in the leaf, so
+                # they compare on the 2.11 line only.
+                $_pinNeeds211 = $false
+                $_pinRocmVer = $null
+                if ($_pinLeaf -like 'gfx*') {
+                    $_pinNeeds211 = $true
+                } elseif ($_pinLeaf -match '^rocm(\d+)\.(\d+)') {
+                    $_pinRocmVer = "$($Matches[1]).$($Matches[2])"
+                    $_pinNeeds211 = ([int]$Matches[1] -gt 7) -or ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -ge 2)
+                }
+                $_instRocmVer = $null
+                if ($torchVer -match '\+rocm(\d+)\.(\d+)') { $_instRocmVer = "$($Matches[1]).$($Matches[2])" }
+                $_instIs211 = $false
+                if ($torchVer -match '^(\d+)\.(\d+)') {
+                    $_instIs211 = ([int]$Matches[1] -gt 2) -or ([int]$Matches[1] -eq 2 -and [int]$Matches[2] -ge 11)
+                }
+                if ($_pinRocmVer -and $_instRocmVer) {
+                    # Both ROCm versions readable: compare them exactly.
+                    $expectedTorchTag = "rocm$_pinRocmVer"
+                    $installedTorchTag = "rocm$_instRocmVer"
+                } else {
+                    # gfx pin or unreadable version: compare on the torch 2.11 line.
+                    $expectedTorchTag  = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
+                    $installedTorchTag = if ($_instIs211)   { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
+                }
            } elseif ($_pinLeaf -like 'cu*' -or $_pinLeaf -eq 'cpu') {
                $expectedTorchTag = $_pinLeaf
            } else {
@ -2584,6 +2612,13 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
            }
        } elseif ($HasNvidiaSmi) {
            $expectedTorchTag = Get-PytorchCudaTag
+        } elseif ($HasROCm -or $script:ROCmGfxArch) {
+            # AMD/ROCm host with no explicit pin: an existing +rocm wheel is the
+            # correct build, not stale. (gfx arch counts even when $HasROCm is
+            # false -- name-inferred Adrenalin hosts still get ROCm torch below.)
+            # Without this an unpinned ROCm venv compares "rocm" != "cpu" and is
+            # needlessly rebuilt, and an installer-managed setup exits as stale.
+            $expectedTorchTag = "rocm"
        } else {
            $expectedTorchTag = "cpu"
        }
@ -2899,11 +2934,15 @@ if ($ROCmIndexUrl) {
    }
 }

-if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
+if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) {
    substep "installing PyTorch (CPU-only)..."
    # After an AMD ROCm fallback, force-reinstall so a partially-installed ROCm torch
    # (which still satisfies the CPU torch>= range) is replaced by the CPU build. Skip
    # the forced reinstall on a genuine CPU-only host so the common path stays fast.
+    # The $ROCmCpuFallback term matters when a PINNED ROCm index failed: $CuTag is
+    # still the rocm/gfx leaf (not "cpu"), so without it execution would fall through
+    # to the CUDA branch and install from the CPU index WITHOUT --force-reinstall,
+    # leaving the partial ROCm torch in place.
    # Build the array directly: an if-expression collapses @("x") to a scalar string,
    # which @splat would then enumerate char-by-char into broken single-letter args.
    $cpuForce = @()