install: honor pinned torch index over CVD/GPU gates and fix leaf-based ROCm classification

- install_python_stack.py: an explicit cu* pin now clears the CUDA_VISIBLE_DEVICES
  empty/-1 hide gate as well as the NVIDIA-presence gate, so
  CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update repairs to CUDA wheels
  (parity with install.sh's get_torch_index_url override, which skips all GPU
  probing). Unpinned CVD=-1 still skips.
- install_python_stack.py: _ensure_cpu_torch installs the bounded _CPU_TORCH_PKG_SPEC
  instead of a bare torch/torchvision/torchaudio trio; the /cpu index now also
  serves torch 2.11+, which is outside the supported <2.11 range.
- install.sh: the torch>=2.11 constraint case matches the index leaf (rocm7.2|gfx*)
  instead of the whole URL, so a mirror base path containing a gfx/rocm7.2 segment
  with a cu*/cpu family is not false-matched onto the 2.11 line.
- setup.ps1: the stale-venv check expects rocm torch only for arches the install
  path maps to a repo.amd.com wheel index; an unmapped/unreadable arch installs
  CPU, so a correct CPU venv is no longer marked stale.
- Tests for each of the above.

install.sh

@@ -2482,8 +2482,13 @@ esac
 # 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*)
+# Match on the FINAL path segment ($_torch_index_leaf, computed above), NOT the
+# whole URL: a custom UNSLOTH_PYTORCH_MIRROR whose base path contains a "gfx" or
+# "rocm7.2" segment (e.g. https://mirror.local/gfx-cache) with a cu*/cpu family
+# must not be treated as an AMD per-arch index and pushed to the 2.11 line. This
+# mirrors the leaf-only backend classification just above.
+case "$_torch_index_leaf" in
+    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"

studio/install_python_stack.py

@@ -115,6 +115,12 @@ _CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = (
     "torchaudio>=2.4,<2.11.0",
 )
 
+# CPU torch repair specs (see _ensure_cpu_torch). Same bounds/reasoning as the
+# CUDA spec above: the /cpu index now also publishes torch 2.11+, so a bare trio
+# from the exclusive --index-url would resolve outside the repo's supported
+# <2.11 range (and can pick a torchvision built against a different torch major).
+_CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC
+
 # torchao's C++ extensions are built against ONE exact torch release; a newer
 # torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions
 # due to incompatible torch version ...") and fall back to slow Python. Because
@@ -1126,18 +1132,23 @@ def _ensure_cuda_torch() -> None:
     # Never undo a deliberate ROCm install (setup.ps1 sets this marker).
     if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
         return
+    # An explicit CUDA wheel-index pin (headless / container / CI cross-install)
+    # commits to CUDA wheels and, like install.sh's get_torch_index_url override,
+    # skips ALL host-GPU probing -- so it must clear BOTH the CUDA_VISIBLE_DEVICES
+    # hide gate and the NVIDIA-presence gate below, not just the latter. Otherwise
+    # `CUDA_VISIBLE_DEVICES=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update`
+    # (the exact GPU-less CI case this override targets) would still bail here.
+    _cuda_pinned = _explicit_cuda_torch_index_url() is not None
     # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU (for
     # example a mixed AMD+NVIDIA host that runs ROCm torch on the AMD card);
-    # never force CUDA wheels over that choice.
+    # never force CUDA wheels over that choice unless a CUDA index is pinned.
     _cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
-    if _cvd is not None and _cvd.strip() in ("", "-1"):
+    if not _cuda_pinned and _cvd is not None and _cvd.strip() in ("", "-1"):
         return
     # Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu()
     # covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent.
-    # An explicit CUDA wheel-index pin (headless / container / CI cross-install)
-    # commits to CUDA wheels regardless of whether a GPU is visible here, so it
-    # overrides the GPU-presence gate.
-    if not _has_usable_nvidia_gpu() and _explicit_cuda_torch_index_url() is None:
+    # The explicit CUDA pin overrides the GPU-presence gate too.
+    if not _cuda_pinned and not _has_usable_nvidia_gpu():
         return
 
     # Classify the installed torch: "hip" (ROCm build -- the poisoning
@@ -1267,14 +1278,18 @@ def _ensure_cpu_torch() -> None:
         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.
+    # Pin to the supported torch<2.11 family (same bounds as the CUDA/ROCm repair
+    # specs). The /cpu index now also serves torch 2.11+, so a bare trio off the
+    # exclusive --index-url could resolve outside the supported range or drag in
+    # an ABI-mismatched torchvision/torchaudio.
+    _torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC
     pip_install(
         "CPU torch repair",
         "--force-reinstall",
         "--no-cache-dir",
-        "torch",
-        "torchvision",
-        "torchaudio",
+        _torch_pkg,
+        _vision_pkg,
+        _audio_pkg,
         "--index-url",
         pin,
         constrain = False,

studio/setup.ps1

@@ -2618,7 +2618,24 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
             # 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"
+            # But only the arches the install path below maps to a repo.amd.com
+            # wheel index (the $archFamilyMap set) actually get ROCm torch; an
+            # unmapped/unknown arch (e.g. name-inferred RDNA 2 gfx103X) or an
+            # unreadable arch -- even with $HasROCm -- falls back to CPU torch
+            # there (see the $archFamily null branch below). For those, expect
+            # "cpu" so a correct CPU venv is not marked stale and rebuilt on every
+            # update (or aborted under installer-managed setup).
+            $_rocmWheelArches = @(
+                "gfx1201", "gfx1200",           # RDNA 4
+                "gfx1151", "gfx1150",           # RDNA 3.5 (Strix Halo/Point)
+                "gfx1103", "gfx1102", "gfx1101", "gfx1100",  # RDNA 3
+                "gfx90a", "gfx908"              # MI200 / MI100
+            )
+            if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) {
+                $expectedTorchTag = "rocm"
+            } else {
+                $expectedTorchTag = "cpu"
+            }
         } else {
             $expectedTorchTag = "cpu"
         }

tests/python/test_cross_platform_parity.py

@@ -210,3 +210,52 @@ class TestTorchIndexOverrideParity:
         assert (
             "TorchIndexPinned" in text
         ), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
+
+    def test_cuda_pin_overrides_cvd_hide_gate(self):
+        # A pinned cu* index skips ALL host-GPU probing (parity with install.sh's
+        # get_torch_index_url override), so the Python CUDA repair must let the pin
+        # clear the CUDA_VISIBLE_DEVICES hide gate, not just the NVIDIA-presence
+        # gate. Otherwise CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update
+        # (the GPU-less CI case) would bail before repairing.
+        text = STACK_PY.read_text(encoding = "utf-8")
+        m = re.search(
+            r"def _ensure_cuda_torch\(\).*?(?=\ndef )", text, re.DOTALL
+        )
+        assert m, "could not locate _ensure_cuda_torch"
+        body = m.group(0)
+        # The CVD hide-gate return must be guarded by the CUDA-pin flag.
+        assert "_cuda_pinned" in body, (
+            "_ensure_cuda_torch should compute a CUDA-pin flag so the pin can "
+            "override the CVD hide gate"
+        )
+        assert re.search(
+            r'if not _cuda_pinned and _cvd is not None', body
+        ), "the CVD hide gate must be bypassed when a CUDA index is pinned"
+
+    def test_cpu_repair_pins_supported_torch_range(self):
+        # The explicit-CPU repair must not install a bare torch trio: the /cpu
+        # index now also serves torch 2.11+, so a bare install off the exclusive
+        # --index-url can resolve outside the repo's supported <2.11 range or pull
+        # an ABI-mismatched companion. It must use the bounded CPU/CUDA spec.
+        text = STACK_PY.read_text(encoding = "utf-8")
+        m = re.search(
+            r"def _ensure_cpu_torch\(\).*?(?=\ndef )", text, re.DOTALL
+        )
+        assert m, "could not locate _ensure_cpu_torch"
+        body = m.group(0)
+        assert "_CPU_TORCH_PKG_SPEC" in body, (
+            "_ensure_cpu_torch should install the bounded _CPU_TORCH_PKG_SPEC, "
+            "not a bare torch/torchvision/torchaudio trio"
+        )
+
+    def test_setup_ps1_stale_check_gates_rocm_on_supported_arch(self):
+        # The stale-venv check must only expect ROCm torch for arches the install
+        # path actually maps to a repo.amd.com wheel index. An unmapped arch
+        # (name-inferred RDNA 2 gfx103X) or an unreadable arch installs CPU torch,
+        # so expecting "rocm" there marks a correct CPU venv stale and rebuilds it
+        # every update (or aborts under installer-managed setup).
+        text = SETUP_PS1.read_text(encoding = "utf-8")
+        assert "_rocmWheelArches" in text, (
+            "setup.ps1 stale check should restrict the ROCm expected-tag to the "
+            "supported gfx wheel arches"
+        )

tests/studio/install/test_cuda_repair.py

@@ -148,6 +148,22 @@ class TestCudaRepairFires:
         assert mock_pip.call_count == 1
         assert "cu128" in _index_url(mock_pip)
 
+    def test_cvd_hidden_but_explicit_cuda_pin_repairs(self):
+        # CUDA_VISIBLE_DEVICES=-1 hides the GPU, but an explicit cu* pin skips ALL
+        # host-GPU probing (like install.sh's get_torch_index_url override), so the
+        # CVD hide gate must not suppress the repair. This is the exact GPU-less CI
+        # case the override targets: CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128.
+        for _cvd in ("-1", ""):
+            mock_pip = _run_cuda_repair(
+                nvidia = False,
+                backend = "cuda",
+                cvd = _cvd,
+                index_family = "cu128",
+                torch_state = "hip",
+            )
+            assert mock_pip.call_count == 1
+            assert "cu128" in _index_url(mock_pip)
+
 
 # No-op cases.
 

tests/studio/install/test_rocm_support.py

@@ -2876,6 +2876,23 @@ class TestStrixRocm71Override:
         source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
         assert "TORCH_CONSTRAINT" in source and "2.11" in source
 
+    def test_torch_constraint_211_matches_leaf_not_whole_url(self):
+        """The 2.11 constraint case must match the index LEAF, not the whole URL.
+
+        A custom UNSLOTH_PYTORCH_MIRROR whose base path contains a gfx/rocm7.2
+        segment (e.g. https://mirror.local/gfx-cache) with a cu*/cpu family must
+        not be pushed to the torch 2.11 line -- same leaf-only reasoning the
+        UNSLOTH_TORCH_BACKEND classification uses.
+        """
+        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+        # The 2.11 constraint block must switch on $_torch_index_leaf, not on the
+        # full $TORCH_INDEX_URL (which the earlier, buggy version matched with
+        # */gfx* and would false-positive on a mirror base path).
+        assert 'case "$_torch_index_leaf" in\n    rocm7.2|gfx*)' in source, (
+            "the torch>=2.11 constraint must match the index leaf (rocm7.2|gfx*), "
+            "not the whole URL"
+        )
+
     def test_amd_rocm_mirror_env_var_respected(self):
         """install.sh must honour UNSLOTH_AMD_ROCM_MIRROR for air-gapped installs."""
         source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")