diff --git a/install.sh b/install.sh index 53b09cffab..964a715f22 100755 --- 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 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 0aa6116a08..ee0cd8eafa 100644 --- 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( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index ffbef0bac8..e2a68f90fb 100644 --- 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 = @()