diff --git a/install.ps1 b/install.ps1 index f7f9540970..b895ec24d9 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1939,6 +1939,17 @@ exit 0 # Mirrors Get-PytorchCudaTag in setup.ps1. function Get-TorchIndexUrl { $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } + # Explicit pin -- skip ALL GPU probing when the caller names the wheel index + # (headless / CI / cross-install). Matches install.sh::get_torch_index_url and + # install_python_stack.py: UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); + # UNSLOTH_TORCH_INDEX_FAMILY is the convenience leaf (cpu, cu128, rocm6.4, ...) + # appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured. + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { + return $env:UNSLOTH_TORCH_INDEX_URL.Trim().TrimEnd('/') + } + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) { + return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))" + } if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe @@ -2016,6 +2027,11 @@ exit 0 } catch { return $null } } + # An explicit UNSLOTH_TORCH_INDEX_URL / _FAMILY pin is authoritative: the AMD + # ROCm reroute below must not rewrite it (e.g. a deliberate cpu pin on an AMD + # host, or a pinned ROCm family we already resolved in Get-TorchIndexUrl). + $TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or ` + (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) $TorchIndexUrl = Get-TorchIndexUrl # ── GPU arch → newest compatible Windows ROCm wheel release ── @@ -2027,7 +2043,9 @@ exit 0 # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. $ROCmIndexUrl = $null $ROCmTorchFloor = $null - if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { + $PinnedRocmVisionSpec = $null + $PinnedRocmAudioSpec = $null + if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 @@ -2077,6 +2095,30 @@ exit 0 } } + # An explicit gfx*/rocm pin skips the auto-reroute above, but the generic + # CPU/CUDA install below would use torch>=2.4,<2.11 and pull a known-bad wheel + # on the gfx115x/gfx120x/rocm>=7.2 indexes (the torch._C._grouped_mm null-ptr + # bug). Route a pinned ROCm index through the ROCm install path with the same + # 2.11 floor/companions the unpinned reroute derives from the gfx arch. + if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) { + $_pinLeaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLower() + $_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 = $TorchIndexUrl + $ROCmTorchFloor = "torch>=2.11.0,<2.12.0" + $PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0" + $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0" + substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan" + } elseif ($_pinLeaf -like 'rocm*') { + # Older rocm (<=7.1) ships torch <2.11; route via the ROCm path with the + # default floor so the pinned family resolves its own wheels. + $ROCmIndexUrl = $TorchIndexUrl + } + } + if ($ROCmIndexUrl) { $TorchIndexFamily = "rocm" } else { @@ -2189,18 +2231,22 @@ exit 0 $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin the companions to match $torchSpec; bare names can resolve an # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { # Transient AMD-index failure: fall back to a CPU base so the install # still completes; Studio setup retries ROCm afterwards. + # 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" } substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Studio setup retries ROCm." "Yellow" # --force-reinstall: a failed ROCm install can leave an unpinned ROCm # 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 } + $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) @@ -2310,8 +2356,8 @@ exit 0 $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin companions like the fresh ROCm path (bare names can pull an # ABI-incompatible torchvision/torchaudio from the per-arch index). - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { diff --git a/install.sh b/install.sh index 548e6f702a..53b09cffab 100755 --- a/install.sh +++ b/install.sh @@ -1985,6 +1985,28 @@ _has_usable_nvidia_gpu() { get_torch_index_url() { _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" _base="${_base%/}" + # Explicit override -- skip ALL GPU probing when the caller pins the wheel + # index. Headless / container / CI builds (and anyone cross-installing for a + # different target) must not let the build host's GPU -- or the lack of one -- + # decide the wheel family. This is the same "tell the build, don't ask the + # hardware" approach the Docker base image and vLLM/SGLang's Dockerfiles take. + # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); UNSLOTH_TORCH_INDEX_FAMILY + # is the convenience form (cpu, cu124, cu126, cu128, cu130, rocm6.4, ...) + # appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured. + if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ]; then + # Strip ALL trailing slashes (match the Python side's .rstrip("/") and the + # Strix mirror handling below) -- a double/triple-slash URL 404s on strict + # pip proxies (artifactory, sonatype). + _url="${UNSLOTH_TORCH_INDEX_URL}" + while [ "${_url%/}" != "$_url" ]; do _url="${_url%/}"; done + echo "$_url"; return + fi + if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then + _family="${UNSLOTH_TORCH_INDEX_FAMILY}" + while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done + while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done + echo "$_base/$_family"; return + fi # macOS: always CPU (no CUDA support) case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac # Try nvidia-smi -- require the binary to actually list a usable GPU. @@ -2381,7 +2403,16 @@ _maybe_bootstrap_rocm_wsl() { [ -n "$_rw_tmp" ] && rm -f "$_rw_tmp" return 0 } -_maybe_bootstrap_rocm_wsl || true +# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY), +# honour it everywhere downstream: skip the WSL ROCm bootstrap (which can run +# sudo + large downloads after probing /dev/dxg) and the Radeon/Strix rerouting +# below (which would re-probe the GPU and overwrite the pinned URL). A headless / +# container / CI build must get exactly the index it asked for. +_torch_index_pinned=false +if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] || [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then + _torch_index_pinned=true +fi +[ "$_torch_index_pinned" = true ] || _maybe_bootstrap_rocm_wsl || true TORCH_INDEX_URL=$(get_torch_index_url) @@ -2400,16 +2431,23 @@ case "$_torch_index_leaf" in *) export UNSLOTH_TORCH_BACKEND="cuda" ;; esac -# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. +# rocm7.2 and the AMD per-gfx indexes (repo.amd.com/.../gfxNNNN) ship torch +# 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. # All other ROCm tags and CUDA stay within <2.11.0. case "$TORCH_INDEX_URL" in - */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; + */rocm7.2|*/gfx*) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; esac # Auto-detect GPU for AMD ROCm based # get_torch_index_url must have chosen */rocm* # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". +# Skipped entirely when the index is pinned: an explicit override (even a ROCm +# one like UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4) must not be rerouted to the +# Radeon/Strix repos by GPU probing. _amd_gpu_radeon=false +if [ "$_torch_index_pinned" = false ]; then case "$TORCH_INDEX_URL" in */rocm*) if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \ @@ -2490,6 +2528,7 @@ case "$TORCH_INDEX_URL" in fi ;; esac +fi # _torch_index_pinned guard (Radeon + Strix reroute) _TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL") if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then _TAURI_TORCH_INDEX_FAMILY="radeon" diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 9060b57542..0aa6116a08 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -933,11 +933,20 @@ def _detect_cuda_torch_index_url() -> str: Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update` repairs to the same wheel family a fresh `curl | sh` install would pick. - Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the - legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings. - Defaults to cu126 when nvidia-smi is missing or the version is unreadable - (e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback). + Honours the same explicit overrides first (UNSLOTH_TORCH_INDEX_URL / + UNSLOTH_TORCH_INDEX_FAMILY) so a headless / container / CI install never lets + the host GPU decide the wheel family. Otherwise probes nvidia-smi (PATH, then + /usr/bin/nvidia-smi) and parses both the legacy "CUDA Version:" and the newer + "CUDA UMD Version:" spellings. Defaults to cu126 when nvidia-smi is missing or + the version is unreadable (e.g. NVIDIA detected only via the + /proc/driver/nvidia/gpus fallback). """ + _override_url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + if _override_url: + return _override_url.rstrip("/") + _override_family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + if _override_family: + return f"{_PYTORCH_WHL_BASE}/{_override_family.strip('/')}" exe = shutil.which("nvidia-smi") if not exe and os.path.isfile("/usr/bin/nvidia-smi"): exe = "/usr/bin/nvidia-smi" @@ -972,6 +981,30 @@ def _detect_cuda_torch_index_url() -> str: return f"{_PYTORCH_WHL_BASE}/{tag}" +def _explicit_torch_index_url() -> "str | None": + """The wheel index URL pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY, else None. + + Lets the CUDA/ROCm repair helpers honour the exact pinned family/URL instead + of re-probing the GPU. Mirrors install.sh::get_torch_index_url's override. + """ + url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + if url: + return url.rstrip("/") + family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + if family: + return f"{_PYTORCH_WHL_BASE}/{family.strip('/')}" + return None + + +def _explicit_rocm_torch_index_url() -> "str | None": + """The pinned wheel index URL when it names a ROCm family (rocm*/gfx*), else None.""" + url = _explicit_torch_index_url() + if url is None: + return None + leaf = url.rstrip("/").rsplit("/", 1)[-1].lower() + return url if leaf.startswith(("rocm", "gfx")) else None + + def _ensure_cuda_torch() -> None: """Repair a venv whose torch is a ROCm build on an NVIDIA host. @@ -1004,7 +1037,10 @@ def _ensure_cuda_torch() -> None: return # Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu() # covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent. - if not _has_usable_nvidia_gpu(): + # 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_torch_index_url() is None: return # Classify the installed torch: "hip" (ROCm build -- the poisoning @@ -1017,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, @@ -1037,16 +1075,35 @@ def _ensure_cuda_torch() -> None: _marker_lines = [ line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() ] - if not _marker_lines or _marker_lines[-1] != "hip": - return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is + if not _marker_lines: + return + _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 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() + _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 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 print( - f" torch is a ROCm build on an NVIDIA host -- reinstalling " - f"CUDA torch from {index_url}\n" - f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch " - f"on a mixed AMD+NVIDIA host)" + f" {_why} -- reinstalling CUDA torch from {index_url}\n" + f" (set UNSLOTH_TORCH_BACKEND=rocm or cpu to keep a deliberate " + f"non-CUDA torch)" ) pip_install( "CUDA torch repair", @@ -1247,7 +1304,9 @@ def _ensure_rocm_torch() -> None: # an incompatible wheel. Use HIP_VISIBLE_DEVICES for the runtime target. _strix_override_url: "str | None" = None _strix_override_pkgs: "tuple[str, str, str] | None" = None - if ver < (7, 2): + # An explicit ROCm wheel-index pin is authoritative: never auto-reroute it to + # the AMD per-gfx index (the caller already chose the family/URL). + if ver < (7, 2) and _explicit_rocm_torch_index_url() is None: gfx_codes = _detect_amd_gfx_codes() _strix_gfx = {"gfx1151", "gfx1150"} _detected_strix = _strix_gfx.intersection(gfx_codes) @@ -1310,23 +1369,35 @@ def _ensure_rocm_torch() -> None: ) rocm_torch_ready = True elif not has_hip_torch: - # Select best matching wheel tag (newest ROCm version <= installed) - tag = next( - ( - t - for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) - if ver >= (maj, mn) - ), - None, - ) + # Honour an explicit ROCm wheel-index pin verbatim instead of re-detecting + # the host ROCm version; otherwise select the best wheel tag (newest ROCm + # version <= installed). gfx*/rocm7.2 indexes serve torch 2.11+, so match + # the constraints to the index leaf when overridden. + _override_idx = _explicit_rocm_torch_index_url() + if _override_idx is not None: + index_url = _override_idx + tag = index_url.rstrip("/").rsplit("/", 1)[-1].lower() + else: + tag = next( + ( + t + for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) + if ver >= (maj, mn) + ), + None, + ) if tag is None: print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall") else: - index_url = f"{_PYTORCH_WHL_BASE}/{tag}" - print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}") - _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( - tag, _ROCM_TORCH_PKG_SPECS["_default"] - ) + if _override_idx is None: + index_url = f"{_PYTORCH_WHL_BASE}/{tag}" + print(f" ROCm torch -- installing from {index_url}") + if tag.startswith("gfx"): + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["rocm7.2"] + else: + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( + tag, _ROCM_TORCH_PKG_SPECS["_default"] + ) pip_install( f"ROCm torch ({tag})", "--force-reinstall", @@ -1420,6 +1491,23 @@ def _infer_no_torch() -> bool: # GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown # (standalone `unsloth studio update` runs, where we re-detect normally). _TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower() +# When install.sh did not run (standalone `unsloth studio update`) but the caller +# pinned the wheel index explicitly, derive the backend from that override so the +# CUDA/ROCm repair helpers honour it instead of re-probing the GPU and possibly +# reinstalling a different family. Classify on the final URL/family segment, +# mirroring install.sh's UNSLOTH_TORCH_BACKEND case. +if not _TORCH_BACKEND: + _idx_override = ( + os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + or os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + ) + _idx_leaf = _idx_override.rstrip("/").rsplit("/", 1)[-1].lower() + if _idx_leaf.startswith(("rocm", "gfx")): + _TORCH_BACKEND = "rocm" + elif _idx_leaf == "cpu": + _TORCH_BACKEND = "cpu" + elif _idx_leaf.startswith("cu"): + _TORCH_BACKEND = "cuda" def _torch_step_label(suffix: str) -> str: diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 38398eb5f5..ffbef0bac8 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -402,6 +402,29 @@ function Get-PytorchCudaTag { return "cu126" } +# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the +# stale-venv check and the install selection below so a pinned wheel index wins +# over GPU probing -- matching install.sh, install.ps1 and install_python_stack.py. +# UNSLOTH_TORCH_INDEX_URL is verbatim (full URL); _FAMILY is the leaf (cpu, cu128, +# rocm6.4, ...) joined to the mirror base so UNSLOTH_PYTORCH_MIRROR is honoured. +function Get-PinnedTorchIndexUrl { + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { + return $env:UNSLOTH_TORCH_INDEX_URL.Trim().TrimEnd('/') + } + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) { + $base = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } + return "$base/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))" + } + return $null +} + +# The last path segment of a wheel index URL (cu128 / cpu / rocm6.4 / gfx1151). +function Get-TorchIndexLeaf { + param([string]$Url) + if ([string]::IsNullOrWhiteSpace($Url)) { return $null } + return ($Url.TrimEnd('/') -split '/')[-1].ToLowerInvariant() +} + # VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major # (18->v180, 17->v170), defaulting to v170 when unparseable. function Get-VcBuildCustomizationsDir { @@ -2517,6 +2540,11 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode if ($finished -and $proc.ExitCode -eq 0 -and $torchVer) { if ($torchVer -match '\+(cu\d+)') { $installedTorchTag = $Matches[1] + } elseif ($torchVer -match '\+rocm') { + # Any +rocm / gfx wheel -> generic "rocm" flavor. The exact ROCm + # version is repaired later by install_python_stack.py; here we + # only need the flavor so a correct ROCm venv is not marked stale. + $installedTorchTag = "rocm" } elseif ($torchVer -match '\+cpu') { $installedTorchTag = "cpu" } else { @@ -2536,8 +2564,30 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode } if (-not $shouldRebuild) { - $expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" } - if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { + $_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. + 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 ($_expectedKnown -and $installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { $shouldRebuild = $true } } @@ -2712,7 +2762,13 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User') substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" -if ($HasNvidiaSmi) { +# Explicit pin (URL or family) wins over GPU probing and suppresses the AMD +# reroute below; matches install.sh / install.ps1 / install_python_stack.py. +$PinnedTorchIndexUrl = Get-PinnedTorchIndexUrl +$TorchIndexPinned = [bool]$PinnedTorchIndexUrl +if ($PinnedTorchIndexUrl) { + $CuTag = Get-TorchIndexLeaf $PinnedTorchIndexUrl +} elseif ($HasNvidiaSmi) { $CuTag = Get-PytorchCudaTag } else { $CuTag = "cpu" @@ -2733,7 +2789,7 @@ $ROCmIndexUrl = $null # SDK -- which flips Studio out of chat-only (CHAT_ONLY) and enables Train/Export. # Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed # ROCm install still falls back to CPU below, so this is safe. -if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") { +if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 @@ -2783,8 +2839,40 @@ if (($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. +# 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) { substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..." @@ -2821,11 +2909,11 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") { $cpuForce = @() if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") } if ($script:UnslothVerbose) { - Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" + Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl $torchInstallExit = $LASTEXITCODE $output = "" } else { - $output = Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" | Out-String + $output = Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl | Out-String $torchInstallExit = $LASTEXITCODE } if ($torchInstallExit -ne 0) { @@ -2837,11 +2925,11 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") { substep "installing PyTorch with CUDA support ($CuTag)..." substep "(This download is ~2.8 GB -- may take a few minutes)" if ($script:UnslothVerbose) { - Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" + Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl $torchInstallExit = $LASTEXITCODE $output = "" } else { - $output = Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" | Out-String + $output = Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl | Out-String $torchInstallExit = $LASTEXITCODE } if ($torchInstallExit -ne 0) { diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 666ea7ce10..ae1e247634 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -10,6 +10,8 @@ REPO_ROOT = Path(__file__).resolve().parents[2] INSTALL_SH = REPO_ROOT / "install.sh" INSTALL_PS1 = REPO_ROOT / "install.ps1" +SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1" +STACK_PY = REPO_ROOT / "studio" / "install_python_stack.py" class TestNoTorchBackendAutoInInstallSh: @@ -180,3 +182,31 @@ def test_install_ps1_preserves_timeout_override(self): assert ( '$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text ), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT" + + +class TestTorchIndexOverrideParity: + """Every installer must honor UNSLOTH_TORCH_INDEX_URL / _FAMILY so a pinned wheel + index wins over GPU probing on all platforms (no asymmetric, per-OS coverage).""" + + @pytest.mark.parametrize( + "path", + [INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY], + ids = ["install.sh", "install.ps1", "setup.ps1", "install_python_stack.py"], + ) + def test_installer_reads_override_env(self, path): + text = path.read_text(encoding = "utf-8") + for var in ("UNSLOTH_TORCH_INDEX_URL", "UNSLOTH_TORCH_INDEX_FAMILY"): + assert var in text, f"{path.name} does not honor {var}" + + @pytest.mark.parametrize( + "path", + [INSTALL_PS1, SETUP_PS1], + ids = ["install.ps1", "setup.ps1"], + ) + def test_amd_reroute_guarded_when_pinned(self, path): + # The AMD ROCm reroute must be skipped when the index is explicitly pinned, + # so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten. + text = path.read_text(encoding = "utf-8") + assert ( + "TorchIndexPinned" in text + ), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag" diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 6656142625..0c538f0a12 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -379,6 +379,50 @@ _result=$(run_func "$_dir" " -1 ") assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" rm -rf "$_dir" +# --- explicit overrides (headless / container / CI; no GPU probing) ---------- +# 39) UNSLOTH_TORCH_INDEX_FAMILY pins the family with no GPU present -> that +# family (not the cpu fallback detection would pick). +_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none") +assert_eq "family override (no GPU) -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" + +# 40) Family override wins over a real detection: a host whose nvidia-smi reports +# 12.6 still gets cu128. This is the exact Docker-build case (the builder sees +# the host driver but must publish a cu128 image). +_dir=$(make_mock_smi "12.6") +_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "$_dir") +assert_eq "family override beats detected 12.6 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" +rm -rf "$_dir" + +# 41) UNSLOTH_TORCH_INDEX_URL is used verbatim and wins over detection. +_dir=$(make_mock_smi "12.6") +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu999" run_func "$_dir") +assert_eq "url override beats detection -> verbatim" "https://mirror.example.com/whl/cu999" "$_result" +rm -rf "$_dir" + +# 42) Family override is appended to UNSLOTH_PYTORCH_MIRROR (mirror still honoured). +_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none") +assert_eq "mirror + family override -> mirror/cu128" "https://mirror.example.com/whl/cu128" "$_result" + +# 43) Trailing slash in UNSLOTH_TORCH_INDEX_URL is stripped. +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128/" run_func "none") +assert_eq "url override trailing slash stripped" "https://mirror.example.com/whl/cu128" "$_result" + +# 44) URL override takes precedence over family override. +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu130" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none") +assert_eq "url override beats family override -> url" "https://mirror.example.com/whl/cu130" "$_result" + +# 45) An empty override is ignored (falls through to normal detection). +_result=$(UNSLOTH_TORCH_INDEX_FAMILY="" UNSLOTH_TORCH_INDEX_URL="" run_func "none") +assert_eq "empty overrides ignored -> detected cpu" "https://download.pytorch.org/whl/cpu" "$_result" + +# 46) ALL trailing slashes are stripped from a URL override (not just one). +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128///" run_func "none") +assert_eq "url override double slash stripped" "https://mirror.example.com/whl/cu128" "$_result" + +# 47) Leading and trailing slashes stripped from a family override. +_result=$(UNSLOTH_TORCH_INDEX_FAMILY="//cu128//" run_func "none") +assert_eq "family override slashes stripped" "https://download.pytorch.org/whl/cu128" "$_result" + rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" rm -rf "$_TOOLS_DIR" diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index cea4383268..38e9764ea4 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -64,15 +64,19 @@ def _run_cuda_repair( rocm_marker = False, smi_path = "/usr/bin/nvidia-smi", cvd = None, + index_family = None, ): """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock. - cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it.""" + cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it. + index_family sets UNSLOTH_TORCH_INDEX_FAMILY (the explicit wheel-index pin).""" env = {} if rocm_marker: env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1" if cvd is not None: env["CUDA_VISIBLE_DEVICES"] = cvd + if index_family is not None: + env["UNSLOTH_TORCH_INDEX_FAMILY"] = index_family def _which(name, *a, **k): if name == "nvidia-smi": @@ -99,6 +103,9 @@ def _which(name, *a, **k): stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None) if cvd is None: stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) + if index_family is None: + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None) _ensure_cuda_torch() return mock_pip @@ -128,6 +135,19 @@ def test_rocm_in_version_string_triggers_repair(self): mock_pip = _run_cuda_repair(torch_state = "hip") assert mock_pip.call_count == 1 + def test_no_gpu_but_explicit_cuda_pin_repairs(self): + # Headless / container / CI cross-install: an explicit cu* index pin + # commits to CUDA wheels even though no NVIDIA GPU is visible here, so a + # ROCm-poisoned venv is still repaired (to the pinned family). + mock_pip = _run_cuda_repair( + nvidia = False, + backend = "cuda", + index_family = "cu128", + torch_state = "hip", + ) + assert mock_pip.call_count == 1 + assert "cu128" in _index_url(mock_pip) + # No-op cases. diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 4c50fa1c8e..a900cffb90 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -680,6 +680,36 @@ def test_rocm_72_selects_72_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip) torch_call = mock_pip.call_args_list[0] assert "rocm7.2" in str(torch_call) + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install_try", return_value = True) + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4)) + def test_explicit_gfx_index_honored_and_skips_strix_reroute( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """An explicit gfx wheel-index pin is authoritative: install from it verbatim + with torch 2.11, and never re-probe gfx codes to second-guess it (host ROCm 6.4 + would otherwise pick the rocm6.4 wheel / trigger the Strix re-route).""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall + env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"} + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + # Would raise if the Strix block ran (it is skipped on an explicit pin). + with patch.object( + stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError + ): + _ensure_rocm_torch() + assert mock_pip.call_count == 1 + torch_call = str(mock_pip.call_args_list[0]) + assert "gfx1151" in torch_call + assert "torch>=2.11.0,<2.12.0" in torch_call + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install")