install: honor custom pins and repair pinned venvs in place

Four follow-ups to the torch-index marker work:

- install_python_stack.py: _ensure_cuda_torch/_ensure_rocm_torch now bail when an
  explicit custom-index pin names no known torch family, so a verbatim URL override
  (a private/simple mirror) is not clobbered by auto-detected CUDA/ROCm wheels
  before _ensure_verbatim_torch_index applies it.

- install_python_stack.py: the ROCm marker is additive, not a substitute -- a
  matching marker still runs the family/version check so a wheel swapped after the
  marker was written is caught. Mirrors setup.ps1.

- setup.ps1: a stale venv under an explicit pin, whose torch still imports, is
  repaired in place (force-reinstall torch from the pin in the dependency pass)
  instead of wiped. The wipe path only delegates to install.ps1, so on a direct
  update it stranded the user at "Virtual environment not found" instead of
  applying the new pin. A broken venv or unpinned drift still wipes/delegates.

- install.ps1: when a pinned ROCm install fails over to a CPU base, the marker now
  records the CPU index actually used instead of the ROCm pin, so the next managed
  setup does not see CPU torch under a ROCm pin and abort as stale.

install.ps1

@@ -2080,6 +2080,11 @@ exit 0
     # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
     $ROCmIndexUrl = $null
     $ROCmTorchFloor = $null
+    # Set when a pinned ROCm install fails and a CPU base is installed instead; the
+    # marker must then record the CPU index actually used, not the ROCm pin left in
+    # $TorchIndexUrl (else the next managed setup sees CPU torch under a ROCm pin and
+    # aborts as stale rather than continuing with the intended CPU base/retry path).
+    $RocmCpuFallbackIndexUrl = $null
     $PinnedRocmVisionSpec = $null
     $PinnedRocmAudioSpec = $null
     if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
@@ -2305,6 +2310,9 @@ exit 0
                 # reinstalls ROCm afterwards (recomputes its own index URL).
                 $ROCmIndexUrl = $null
                 $ROCmTorchFloor = $null
+                # Record the CPU index actually installed from, so the marker below
+                # reflects the CPU base rather than the ROCm pin still in $TorchIndexUrl.
+                $RocmCpuFallbackIndexUrl = $CpuFallbackIndexUrl
             }
         } else {
             Write-TauriLog "STEP" "Installing PyTorch"
@@ -2439,10 +2447,11 @@ exit 0
     # Torch is now resolved; write the exact --index-url used so setup.ps1 /
     # install_python_stack.py can detect a later pin change by an EXACT string
     # compare instead of the version-tag heuristic. Reflects the installed family:
-    # $ROCmIndexUrl when the ROCm path ran, else the CUDA/CPU/pinned $TorchIndexUrl.
+    # $ROCmIndexUrl when the ROCm path ran, the CPU fallback index when a pinned
+    # ROCm install failed over to a CPU base, else the CUDA/CPU/pinned $TorchIndexUrl.
     # Skipped for --no-torch (nothing installed). Matches install.sh / setup.ps1.
     if (-not $SkipTorch) {
-        $MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } else { $TorchIndexUrl }
+        $MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } elseif ($RocmCpuFallbackIndexUrl) { $RocmCpuFallbackIndexUrl } else { $TorchIndexUrl }
         Write-TorchIndexMarker -VenvDir $VenvDir -IndexUrl $MarkerIndexUrl
     }
 

studio/install_python_stack.py

@@ -1414,6 +1414,11 @@ def _ensure_cuda_torch() -> None:
     # (or any unrecognised value) are deliberate and must not be overridden.
     if _TORCH_BACKEND not in ("", "cuda"):
         return
+    # An explicit custom-index pin whose leaf names no known torch family wins
+    # VERBATIM (_ensure_verbatim_torch_index applies it); do not override it with
+    # auto-detected CUDA wheels here.
+    if _explicit_unknown_family_torch_index_url() is not None:
+        return
     # No CUDA torch on macOS; Windows venv/torch lifecycle is owned by
     # install.ps1 (and the KFD poisoning bug is Linux-only), so skip both.
     if IS_MACOS or IS_WINDOWS or NO_TORCH:
@@ -1611,6 +1616,14 @@ def _ensure_rocm_torch() -> None:
     # different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.).
     if _TORCH_BACKEND in ("cuda", "cpu"):
         return
+    # An explicit custom-index pin whose leaf names no known torch family (a private
+    # PEP 503 mirror, /simple, /current, ...) wins VERBATIM -- _ensure_verbatim_torch_index()
+    # applies it. Never override such a pin with the auto-detected ROCm index here: the
+    # user chose that index (a valid cross-install / custom-mirror case), and installing
+    # ROCm wheels over it before the verbatim pass runs would leave the URL override
+    # unhonored (the marker then matches the pin, so verbatim skips the restore).
+    if _explicit_unknown_family_torch_index_url() is not None:
+        return
     # setup.ps1 sets this after installing AMD wheels; skip the probe only when
     # torch is actually importable as ROCm. If the venv was wiped between runs,
     # the stale env-var would suppress a needed reinstall.
@@ -1823,10 +1836,19 @@ def _ensure_rocm_torch() -> None:
     _rocm_pin_mismatch = False
     if has_hip_torch and _rocm_pin is not None:
         _marker_verdict = _marker_pin_mismatch(_rocm_pin)
-        if _marker_verdict is None:
+        if _marker_verdict is True:
-            _rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
+            # Marker records a DIFFERENT index than the pin -> reinstall. This is the
+            # only signal that catches a per-arch gfx switch (gfx1151 -> gfx120X-all,
+            # both +rocm7.13.0) the version-tag heuristic below cannot see.
+            _rocm_pin_mismatch = True
         else:
-            _rocm_pin_mismatch = _marker_verdict
+            # Marker matches OR is absent: the marker is an ADDITIONAL rebuild signal,
+            # not a substitute for validating the installed wheel. Still run the
+            # family/version check so a stale wheel (torch swapped after the marker was
+            # written -- e.g. marker records gfx1151 but the venv now carries generic
+            # +rocm7.2 or an older +rocm6.4) is caught. Mirrors setup.ps1, which keeps
+            # the flavor check alongside the marker for this reason.
+            _rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
 
     rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch
 

studio/setup.ps1

@@ -2700,6 +2700,9 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
     $VenvPyExe = Join-Path $VenvDir "Scripts\python.exe"
     $installedTorchTag = $null
     $shouldRebuild = $false
+    # Set when a stale venv under an explicit torch-index pin is repaired in place
+    # (force-reinstall torch from the pin) instead of wiped -- see the rebuild block.
+    $script:PinChangedForceReinstall = $false
 
     if (Test-Path -LiteralPath $VenvPyExe) {
         try {
@@ -2830,6 +2833,19 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
         }
     }
 
+    # A stale venv under an explicit torch-index pin, whose torch still imports, is
+    # repaired IN PLACE: the dependency pass below force-reinstalls torch from the pin
+    # instead of wiping. setup.ps1's rebuild path only removes the venv and delegates
+    # to install.ps1, so on a direct `unsloth studio update` a wipe would strand the
+    # user at "Virtual environment not found" rather than applying the new pin. Only a
+    # broken venv (torch unimportable) or an unpinned GPU-detected drift falls through
+    # to the wipe/delegate path below.
+    if ($shouldRebuild -and $_pinnedIdx -and $installedTorchTag) {
+        substep "Torch-index pin changed ($installedTorchTag) -- reinstalling torch from the pin in place." "Cyan"
+        $script:PinChangedForceReinstall = $true
+        $shouldRebuild = $false
+    }
+
     if ($shouldRebuild) {
         $reason = if ($installedTorchTag) { "torch $installedTorchTag != required $expectedTorchTag" } else { "torch could not be imported" }
         if ($InstallerManagedSetup) {
@@ -2995,6 +3011,11 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1)
 #     pip install unsloth 2>&1 | Out-Null
 # }
 
+# A torch-index pin change repairs in place: force the dependency pass so the torch
+# install below force-reinstalls from the new pin (the fast "up to date" path would
+# otherwise skip it and leave the old wheel under the changed pin).
+if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false }
+
 if (-not $SkipPythonDeps) {
 
 if ($script:UnslothVerbose) {
@@ -3184,7 +3205,9 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) {
     # 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 = @()
-    if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }
+    # --force-reinstall also on a pin change: a stale +cu / +rocm wheel still
+    # satisfies the CPU torch>= range, so uv would keep it and only swap companions.
+    if ($ROCmCpuFallback -or $script:PinChangedForceReinstall) { $cpuForce = @("--force-reinstall") }
     if ($script:UnslothVerbose) {
         Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl
         $torchInstallExit = $LASTEXITCODE
@@ -3201,12 +3224,17 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) {
 } elseif (-not $ROCmIndexUrl) {
     substep "installing PyTorch with CUDA support ($CuTag)..."
     substep "(This download is ~2.8 GB -- may take a few minutes)"
+    # --force-reinstall on a pin change: an installed cuXXX wheel satisfies the bare
+    # torch requirement (PEP 440 ignores the +cuXXX local tag), so without it uv would
+    # keep the old build and the changed CUDA pin (e.g. cu126 -> cu128) never applies.
+    $cudaForce = @()
+    if ($script:PinChangedForceReinstall) { $cudaForce = @("--force-reinstall") }
     if ($script:UnslothVerbose) {
-        Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl
+        Fast-Install torch torchvision torchaudio @cudaForce --index-url $TorchInstallIndexUrl
         $torchInstallExit = $LASTEXITCODE
         $output = ""
     } else {
-        $output = Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl | Out-String
+        $output = Fast-Install torch torchvision torchaudio @cudaForce --index-url $TorchInstallIndexUrl | Out-String
         $torchInstallExit = $LASTEXITCODE
     }
     if ($torchInstallExit -ne 0) {