install: let UNSLOTH_TORCH_INDEX_FAMILY / _URL override CUDA wheel detection#6692
install: let UNSLOTH_TORCH_INDEX_FAMILY / _URL override CUDA wheel detection#6692danielhanchen wants to merge 30 commits into
Conversation
…tection
get_torch_index_url (and the studio-update mirror _detect_cuda_torch_index_url)
chose the torch wheel family solely by probing the host GPU, with no override.
In a headless / container / CI build the host driver is visible via the
/proc/driver/nvidia/gpus fallback but nvidia-smi cannot report a CUDA version,
so the function fell back to its cu126 default and installed the wrong wheels
(e.g. a cu128 image got cu126 torch).
Add an explicit override checked before any probing, in both the shell installer
and the Python studio-update path:
- UNSLOTH_TORCH_INDEX_URL full index URL, used verbatim (wins)
- UNSLOTH_TORCH_INDEX_FAMILY family (cpu, cu128, rocm6.4, ...) appended to the
mirror base (UNSLOTH_PYTORCH_MIRROR still honoured)
This matches how the published GPU images select CUDA -- vLLM and SGLang take the
CUDA version from an explicit build ARG rather than detecting it, and the Unsloth
Docker base image already pins the cu128 index directly. Desktop installs are
unchanged: with no override set, detection runs exactly as before.
Adds test_get_torch_index_url.sh cases for the override (family, full URL,
precedence, mirror base, trailing-slash strip, empty-ignored).
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for overriding the PyTorch wheel index URL and family via the UNSLOTH_TORCH_INDEX_URL and UNSLOTH_TORCH_INDEX_FAMILY environment variables in both the shell installer (install.sh) and the Python updater (install_python_stack.py). This allows headless, container, or CI environments to bypass GPU probing and pin a specific wheel family. Unit tests have also been added to verify these overrides. Feedback was provided regarding an inconsistency in how trailing and leading slashes are stripped between the shell and Python implementations, with a suggestion to use a loop in the shell script to robustly strip multiple slashes.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ]; then | ||
| echo "${UNSLOTH_TORCH_INDEX_URL%/}"; return | ||
| fi | ||
| if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then | ||
| _family="${UNSLOTH_TORCH_INDEX_FAMILY#/}" | ||
| echo "$_base/${_family%/}"; return | ||
| fi |
There was a problem hiding this comment.
The Python implementation in install_python_stack.py uses .rstrip('/') and .strip('/') to strip all leading and trailing slashes from the override variables. However, the shell implementation in install.sh currently only strips a single leading or trailing slash (using ${VAR%/} and ${VAR#/}).
To ensure consistency between the shell installer and the Python updater, and to prevent potential 404 errors on strict pip proxies (as noted in the comment on line 2499), we should strip all leading and trailing slashes in install.sh using a loop.
| if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ]; then | |
| echo "${UNSLOTH_TORCH_INDEX_URL%/}"; return | |
| fi | |
| if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then | |
| _family="${UNSLOTH_TORCH_INDEX_FAMILY#/}" | |
| echo "$_base/${_family%/}"; return | |
| fi | |
| if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ]; then | |
| _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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 138c0949b4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ]; then | ||
| echo "${UNSLOTH_TORCH_INDEX_URL%/}"; return | ||
| fi | ||
| if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then |
There was a problem hiding this comment.
When either override is set, this branch is still reached only after _maybe_bootstrap_rocm_wsl has already run (_maybe_bootstrap_rocm_wsl || true precedes TORCH_INDEX_URL=$(get_torch_index_url)). On a WSL Strix Halo host with no ROCm runtime, a headless build that pins CPU/CUDA via UNSLOTH_TORCH_INDEX_URL or UNSLOTH_TORCH_INDEX_FAMILY can still probe /dev/dxg and run the ROCm bootstrap with sudo/large downloads before the override is honored. Please short-circuit that bootstrap when an explicit index override is present, or resolve the override before calling it.
Useful? React with 👍 / 👎.
| _family="${UNSLOTH_TORCH_INDEX_FAMILY#/}" | ||
| echo "$_base/${_family%/}"; return |
There was a problem hiding this comment.
When UNSLOTH_TORCH_INDEX_FAMILY is set to a ROCm tag on a Radeon host, this returns the pinned mirror URL, but the later _amd_gpu_radeon fresh-install path still probes the GPU and installs from repo.radeon.com instead of $TORCH_INDEX_URL (falling back only on failure). That means air-gapped/CI installs using UNSLOTH_PYTORCH_MIRROR=... UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4 won't use the requested mirror or pinned family whenever ROCm detection succeeds; guard the Radeon/Strix reroutes when an explicit index override was supplied.
Useful? React with 👍 / 👎.
| _override_family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() | ||
| if _override_family: | ||
| return f"{_PYTORCH_WHL_BASE}/{_override_family.strip('/')}" |
There was a problem hiding this comment.
In standalone unsloth studio update, this override is only read inside the CUDA repair helper, while _TORCH_BACKEND remains empty and _ensure_rocm_torch() still runs afterward. On an AMD Linux host with CPU torch, setting UNSLOTH_TORCH_INDEX_FAMILY=cpu (or a CUDA family for a headless build) therefore still lets ROCm GPU detection reinstall ROCm torch from the detected ROCm index, ignoring the requested family; derive the backend from the override before both CUDA and ROCm repair paths.
Useful? React with 👍 / 👎.
Address review feedback on the override added in this PR so a pinned index is honoured everywhere, not just in get_torch_index_url: - Skip the WSL ROCm bootstrap (root privilege + large downloads, probes /dev/dxg) when UNSLOTH_TORCH_INDEX_URL / _FAMILY is set; it previously ran before the override was consulted. - Skip the Radeon/Strix rerouting (which re-probes the GPU and overwrites the resolved URL with repo.radeon.com / repo.amd.com) when the index is pinned, so an explicit ROCm override (e.g. UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4) is kept. - install_python_stack.py: derive _TORCH_BACKEND from the override when UNSLOTH_TORCH_BACKEND is unset (standalone studio update), so _ensure_rocm_torch / _ensure_cuda_torch repair to the requested family instead of re-detecting. - Strip ALL leading/trailing slashes in the shell override to match the Python side (avoids 404s on strict pip proxies). Adds test cases for double-slash and leading/trailing-slash overrides.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b02a609af5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # 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 |
There was a problem hiding this comment.
When a caller pins an AMD per-arch index, for example UNSLOTH_TORCH_INDEX_URL=https://repo.amd.com/rocm/whl/gfx1151/ for a Strix/air-gapped install, this guard skips the Strix block that would have set TORCH_CONSTRAINT to torch>=2.11.0,<2.12.0. The fresh install then reaches the generic PyTorch install with the default <2.11.0 constraint, but these gfx* indexes are the 2.11 ROCm wheel path, so uv cannot select the intended torch build. Please derive the 2.11 constraint from a pinned gfx* leaf before bypassing GPU probing.
Useful? React with 👍 / 👎.
| if _idx_leaf.startswith(("rocm", "gfx")): | ||
| _TORCH_BACKEND = "rocm" |
There was a problem hiding this comment.
The new override handling classifies UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4 or a full .../rocm6.4 URL only as _TORCH_BACKEND = "rocm"; it does not pass the requested family/index into _ensure_rocm_torch(), which still computes its install URL from the detected host ROCm version. In standalone unsloth studio update on an AMD Linux host with CPU torch, pinning rocm6.4 on a ROCm 7.2 machine will therefore reinstall from rocm7.2 instead of the requested family/mirror. Please make the ROCm repair path consume the explicit override, not just the backend classification.
Useful? React with 👍 / 👎.
|
Addressed the review feedback in the latest commit so a pinned index is authoritative everywhere, not just inside
Added test cases for double-slash and leading/trailing-slash overrides; the |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b02a609af5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if _idx_leaf.startswith(("rocm", "gfx")): | ||
| _TORCH_BACKEND = "rocm" |
There was a problem hiding this comment.
When unsloth studio update (or the post-install Python stack step) is run with UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4 or a ROCm UNSLOTH_TORCH_INDEX_URL, this block only collapses the override to _TORCH_BACKEND = "rocm". _ensure_rocm_torch() still re-detects the host ROCm version and installs from f"{_PYTORCH_WHL_BASE}/{tag}" (and can still take the Strix per-gfx override), so a CPU/CUDA torch repair ignores the exact pinned family/URL that the shell path just promised to honor. This breaks the explicit override contract for standalone Studio updates and for installs where the Python repair step runs after a pinned ROCm selection.
Useful? React with 👍 / 👎.
| elif _idx_leaf.startswith("cu"): | ||
| _TORCH_BACKEND = "cuda" |
There was a problem hiding this comment.
For a standalone unsloth studio update in the headless/container case this change classifies UNSLOTH_TORCH_INDEX_FAMILY=cu128 as CUDA, but _ensure_cuda_torch() still returns before using _detect_cuda_torch_index_url() unless _has_usable_nvidia_gpu() succeeds. That means an existing ROCm-poisoned venv on a GPU-less builder is left as ROCm even though the caller explicitly pinned a CUDA wheel family, so the new Python-side override does not actually cover the no-visible-GPU scenario it was added for.
Useful? React with 👍 / 👎.
| # pip proxies (artifactory, sonatype). | ||
| _url="${UNSLOTH_TORCH_INDEX_URL}" | ||
| while [ "${_url%/}" != "$_url" ]; do _url="${_url%/}"; done | ||
| echo "$_url"; return |
There was a problem hiding this comment.
When UNSLOTH_TORCH_INDEX_URL points at an AMD per-gfx index such as https://repo.amd.com/rocm/whl/gfx1151, this return makes TORCH_INDEX_URL end in gfx1151 without taking the Strix reroute path that normally also raises TORCH_CONSTRAINT to torch>=2.11.0,<2.12.0. The only generic constraint adjustment below handles */rocm7.2, so the fresh install still asks that per-gfx index for torch>=2.4,<2.11.0 even though the existing Strix path documents that this index serves the torch 2.11+ ROCm build, causing the pinned full-URL install to fail to resolve.
Useful? React with 👍 / 👎.
Follow-up to the override work in this PR: the get_torch_index_url / install.sh reroute already respect a pinned UNSLOTH_TORCH_INDEX_URL / _FAMILY, but the Python repair helpers in install_python_stack.py still re-probed the GPU and could overwrite the pinned family. Make the pin authoritative there too: - _ensure_cuda_torch: an explicit cu* pin commits to CUDA wheels, so repair a ROCm-poisoned venv even when no NVIDIA GPU is visible here (headless / container / CI cross-install), instead of bailing on the GPU-presence gate. - _ensure_rocm_torch: skip the AMD per-gfx (Strix) reroute when a ROCm index is pinned, and in the generic reinstall path install from the pinned URL verbatim rather than re-detecting the host ROCm version. gfx*/rocm7.2 indexes serve torch 2.11+, so select the 2.11 package specs for a gfx leaf. - install.sh: raise the torch constraint to 2.11 for */gfx* indexes too, matching rocm7.2, so a pinned full-URL/family override that returns early keeps a valid constraint. Add _explicit_torch_index_url / _explicit_rocm_torch_index_url helpers and tests covering the no-GPU CUDA pin repair and the explicit gfx index honored verbatim.
|
Thanks, all three addressed in 997155d.
Added |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0c0e2cbbc7
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| 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: |
There was a problem hiding this comment.
When a ROCm/GFX index is pinned, this guard suppresses the Strix auto-reroute, but the pinned index is only installed later inside the elif not has_hip_torch path. On a Strix ROCm 7.1 venv that already has HIP torch, running UNSLOTH_TORCH_INDEX_URL=https://repo.amd.com/rocm/whl/gfx1151 unsloth studio update now skips the Strix fix and never reinstalls from the pinned index, leaving the known-bad 7.1 wheel in place.
Useful? React with 👍 / 👎.
| _override_idx = _explicit_rocm_torch_index_url() | ||
| if _override_idx is not None: | ||
| index_url = _override_idx |
There was a problem hiding this comment.
The explicit ROCm index is only consulted here, after _ensure_rocm_torch() has already required a visible AMD GPU and a readable host ROCm version. In a headless/container cross-install, or on a builder with only NVIDIA visible, UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4 sets _TORCH_BACKEND to rocm but the function returns before reaching this block, so the pinned ROCm wheel is never installed and the venv keeps its CPU/CUDA torch.
Useful? React with 👍 / 👎.
| # 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" ;; |
There was a problem hiding this comment.
This pattern checks the whole URL, so a custom mirror whose base path contains a gfx* segment, e.g. UNSLOTH_PYTORCH_MIRROR=https://mirror.local/gfx-cache with a CUDA family, is treated like an AMD per-arch index and raises TORCH_CONSTRAINT to torch>=2.11.0,<2.12.0. The backend classification just above intentionally uses only the final path segment to avoid this mirror-path false positive; the constraint should do the same or CUDA/CPU installs from such mirrors can resolve the wrong torch family.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. install.sh's torch>=2.11 constraint case now switches on the index leaf (case $_torch_index_leaf in rocm7.2|gfx*)) instead of the whole URL, so a custom UNSLOTH_PYTORCH_MIRROR whose base path contains a gfx/rocm7.2 segment with a cu*/cpu family is no longer false-matched onto the 2.11 line. This mirrors the leaf-only UNSLOTH_TORCH_BACKEND classification just above, and install.ps1 / setup.ps1 already classify on the leaf.
The pinned-index work landed for install.sh and install_python_stack.py, but the Windows installers still picked the wheel index from GPU probing. Extend the same UNSLOTH_TORCH_INDEX_URL / _FAMILY contract so a pinned index wins on every platform: - install.ps1: Get-TorchIndexUrl returns the pinned URL/family before nvidia-smi probing; the AMD ROCm reroute is skipped when the index is pinned, so an explicit cpu/cu* pin on an AMD host is not overwritten. - studio/setup.ps1: add shared Get-PinnedTorchIndexUrl / Get-TorchIndexLeaf helpers; the stale-venv check, the install selection and the AMD reroute all honor the pin, and the CPU/CUDA install pulls from the resolved index URL. - tests: parity test that all four installers read both override vars and the two Windows installers gate the AMD reroute on the pinned flag.
for more information, see https://pre-commit.ci
|
The 5-reviewer pass found the override was only wired into
The "deleted mainline" items the stale integration diff showed (Colab Cloudflare tunnel hardening, offline export/load helpers, etc.) are commits that landed on Tests: |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5a017ecfea
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| $ROCmIndexUrl = $null | ||
| $ROCmTorchFloor = $null | ||
| if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { | ||
| if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { |
There was a problem hiding this comment.
When a Windows install pins an AMD per-arch index, e.g. UNSLOTH_TORCH_INDEX_URL=https://repo.amd.com/rocm/whl/gfx1151, this new guard skips the only branch that sets $ROCmTorchFloor; the later generic install still runs with "torch>=2.4,<2.11.0" against $TorchIndexUrl. For the gfx115x/gfx120x indexes that the unpinned path already constrains to torch 2.11, this either resolves the known-bad older wheel or fails to select the intended wheel, so the pinned path should derive the same floor/companion constraints from the pinned leaf before bypassing reroute.
Useful? React with 👍 / 👎.
| if (-not $shouldRebuild) { | ||
| $expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" } | ||
| $_pinnedIdx = Get-PinnedTorchIndexUrl | ||
| $expectedTorchTag = if ($_pinnedIdx) { Get-TorchIndexLeaf $_pinnedIdx } elseif ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" } |
There was a problem hiding this comment.
With a pinned ROCm/GFX index, this compares the installed tag to the raw leaf such as gfx1151 or rocm6.4, but the probe above only ever records cuXXX or cpu and treats a +rocm... torch version as cpu. As a result, unsloth studio update on an existing ROCm Windows venv pinned to a gfx/rocm index is always marked stale (or exits in installer-managed mode) even when the installed ROCm torch is correct; normalize these leaves to the same rocm flavor and parse +rocm before comparing.
Useful? React with 👍 / 👎.
| # 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: |
There was a problem hiding this comment.
In standalone unsloth studio update with UNSLOTH_TORCH_INDEX_FAMILY=cu128 on a headless/no-visible-GPU builder, this bypasses the GPU gate, but the helper still returns unless the existing torch probe prints hip; an existing CPU torch is treated as deliberate and never reinstalled from the pinned CUDA index. Since the update path preserves torch rather than preinstalling it from install.sh, the explicit CUDA override remains ineffective for the common CPU-venv-to-CUDA cross-install case.
Useful? React with 👍 / 👎.
Follow-ups to the override work flagged in review: - install.ps1: a pinned gfx*/rocm>=7.2 index previously skipped the AMD reroute that sets the torch>=2.11 floor, so the generic install used torch>=2.4,<2.11 and could resolve the known-bad _grouped_mm wheel. Route a pinned ROCm index through the ROCm install path with the 2.11 floor + companions, and guard the companion-spec lookup so a skipped reroute block cannot null-deref. - studio/setup.ps1: the stale-venv check compared the installed flavor (cuXXX/cpu, with +rocm misread as cpu) against the raw pinned leaf (gfx1151 / rocm6.4), so a correct pinned ROCm venv was always marked stale. Classify +rocm wheels as the generic 'rocm' flavor and normalize a pinned rocm*/gfx* leaf to 'rocm' before comparing (cu* stays specific so cu126-vs-cu128 still rebuilds). - install_python_stack.py: _ensure_cuda_torch now also reinstalls from a pinned CUDA index when the venv carries a CPU wheel (headless CPU-venv-to-CUDA cross-install via 'studio update'), not only when it finds a ROCm build. - tests: parity assertions already cover all four installers honoring the override.
|
Thanks, all three were real gaps in the pinned-index handling. Fixed:
Validation: |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d2d5f90175
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # 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") { |
There was a problem hiding this comment.
When UNSLOTH_TORCH_INDEX_URL or _FAMILY points at a ROCm/GFX index during standalone studio/setup.ps1 updates, this guard skips the only branch that initializes $ROCmIndexUrl, $ROCmTorchSpec, and the pinned torchvision/torchaudio ranges. The script then falls through with $CuTag set to the pinned leaf and installs via the generic CUDA branch using bare torch torchvision torchaudio from that ROCm index, so pinned Strix/RDNA4 indexes can resolve the known-bad <2.11 wheels or ABI-mismatched companions instead of the constrained ROCm triplet used by the unpinned AMD path.
Useful? React with 👍 / 👎.
| # 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. | ||
| $expectedTorchTag = if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { "rocm" } else { $_pinLeaf } |
There was a problem hiding this comment.
When UNSLOTH_TORCH_INDEX_URL is a full custom index whose final path segment is not itself a flavor, e.g. a PEP 503 mirror ending in /simple, this assigns simple as the required torch tag. A fresh Windows install can successfully install cu128 from that URL in install.ps1, then the installer-managed setup.ps1 pass compares cu128 != simple, marks the venv stale, and exits with a repair error; direct updates similarly rebuild a correct environment. Unknown leaves should not be treated as torch flavor tags.
Useful? React with 👍 / 👎.
| else: | ||
| return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is |
There was a problem hiding this comment.
With standalone unsloth studio update on Linux/headless builders, UNSLOTH_TORCH_INDEX_FAMILY=cu128 now bypasses the visible-GPU gate, but the probe collapses every installed CUDA wheel to the generic marker cuda. If the venv already has torch+cu126, this else returns without reinstalling from the pinned cu128 index, so the explicit override still leaves the wrong CUDA family in place; the repair should compare the installed +cuXXX tag when a CUDA family is pinned.
Useful? React with 👍 / 👎.
| $_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 |
There was a problem hiding this comment.
For a pinned ROCm/GFX index, this sets $ROCmIndexUrl to the same URL stored in $TorchIndexUrl; if the ROCm install later fails, the intended “CPU fallback” path still runs uv pip install ... --index-url $TorchIndexUrl, so it retries the failing ROCm mirror instead of using the CPU wheel index. This makes transient or unreachable pinned ROCm mirrors abort the Windows installer rather than leaving a CPU base for Studio to repair later, unlike the unpinned AMD path where $TorchIndexUrl remains */cpu.
Useful? React with 👍 / 👎.
Follow-ups to the previous round: - studio/setup.ps1: a pinned gfx*/rocm>=7.2 index now routes through the ROCm install path with the 2.11 floor + companions (it previously fell through to the CUDA branch with bare torch/torchvision/torchaudio against the ROCm index). The CPU/CUDA fallback index is forced to the CPU wheel index when a ROCm index is active, so a failed pinned-ROCm install does not retry the ROCm mirror. - studio/setup.ps1: the stale-venv check no longer treats an unrecognized pinned URL leaf (e.g. a PEP 503 mirror ending in /simple) as a torch flavor tag, which was marking a correct venv stale; cu*/cpu/rocm/gfx leaves are still compared. - install.ps1: the post-failure CPU fallback uses an explicit CPU index instead of , which for a pinned ROCm index was the ROCm mirror itself (so the 'fallback' just retried the failing index and aborted the installer). - install_python_stack.py: _ensure_cuda_torch now also reinstalls when the venv's CUDA family differs from a pinned one (installed cu126 vs pinned cu128), not only CPU->CUDA; the probe reports the installed cuXXX tag for the comparison.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4a5baba8ce
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| 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 |
There was a problem hiding this comment.
When a pinned rocm*/gfx* index fails, $TorchInstallIndexUrl has already been set to the CPU index, but $CuTag remains the pinned leaf, so execution skips the CPU fallback branch with @cpuForce and lands here instead. These CUDA-branch calls lack --force-reinstall, so uv/pip can keep a partially installed ROCm torch that still satisfies bare torch even though the script warned it was falling back to CPU; reset $CuTag/route through the CPU branch or include the force reinstall for this fallback path.
Useful? React with 👍 / 👎.
| # 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" ;; |
There was a problem hiding this comment.
With a full UNSLOTH_TORCH_INDEX_URL=.../gfx1151/gfx120X override, this new */gfx* case only raises the torch range; the later shell install still passes bare torchvision torchaudio from the same AMD per-arch index. The PowerShell/Python paths in this change pin those companions because AMD's per-gfx indexes publish them independently, so the shell path can resolve a torchvision/torchaudio build for a different torch minor even while torch is constrained to <2.12, leaving the venv ABI-mismatched.
Useful? React with 👍 / 👎.
|
All four were real gaps in the pinned-index handling. Fixed in
Validation: both |
…r window The pinned-ROCm CPU fallback computes an explicit CPU index, but the comment explaining why it cannot reuse $TorchIndexUrl pushed the actual Invoke-InstallCommandRetry / --force-reinstall call more than 600 chars past the "ROCm PyTorch install failed" message, so test_pr5940_followups's window check no longer saw the retry helper. Move the CPU-index computation and its comment above the failure substep so the retrying force-reinstall stays adjacent to the message. No behavior change: same explicit CPU index, same retry, same --force-reinstall.
|
Tested on Windows 11 with an RTX 5060 Ti (driver 595.97, CUDA 13.2), Python 3.13.13. Checked the default detection plus both overrides, including a full reinstall to confirm the override actually changes what gets installed. Default (no override) detects the right family for this card. Fresh local install picked cu130: With URL resolution checked against the actual So Studio also came up cleanly on the branch earlier (UI loaded, GGUF chat worked) on the default cu130 build: Default desktop installs behave exactly as before with nothing set, and the override does what it says for headless/CI/Docker builds. Looks good. |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f6c5e46dad
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| } 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" |
There was a problem hiding this comment.
On a normal Windows AMD/ROCm install with no UNSLOTH_TORCH_INDEX_* pin and no NVIDIA GPU, this new branch records an existing +rocm torch as rocm, but the expected tag below still falls back to cpu. That makes a healthy ROCm venv look stale (torch rocm != required cpu); for installs launched from install.ps1 the managed setup path exits before finishing, and direct studio update rebuilds unnecessarily. Derive the expected tag from $HasROCm/$ROCmGfxArch here or avoid turning unpinned ROCm into a CPU mismatch.
Useful? React with 👍 / 👎.
| 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 |
There was a problem hiding this comment.
When a pinned ROCm index such as UNSLOTH_TORCH_INDEX_FAMILY=rocm7.2 or gfx1151 fails to install, $CuTag remains the ROCm leaf, so execution skips the CPU branch with @cpuForce and reaches this CUDA branch instead. Since $TorchInstallIndexUrl was precomputed as the CPU URL while $ROCmIndexUrl was set, this installs from the CPU index without --force-reinstall, allowing uv to keep a partially installed ROCm torch that already satisfies torch—the exact case the CPU fallback was guarding against.
Useful? React with 👍 / 👎.
| _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() |
There was a problem hiding this comment.
In standalone unsloth studio update with UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4 or a full ROCm URL, this override is only consulted after _ensure_rocm_torch() has already returned for a visible NVIDIA GPU, no visible AMD GPU, or an unreadable ROCm version. That leaves an existing CPU/CUDA torch in place in headless CI/container or mixed-GPU cross-install scenarios, even though the new override is meant to skip host GPU probing like the CUDA pin path does.
Useful? React with 👍 / 👎.
| elif _idx_leaf == "cpu": | ||
| _TORCH_BACKEND = "cpu" |
There was a problem hiding this comment.
When standalone update runs with UNSLOTH_TORCH_INDEX_FAMILY=cpu or a URL ending in /cpu, this sets _TORCH_BACKEND to cpu; both repair helpers then treat that as a skip signal (_ensure_cuda_torch only runs for ""/cuda, and _ensure_rocm_torch returns for cuda/cpu). There is no CPU reinstall path, so an existing CUDA/ROCm torch remains in place even though the explicit CPU pin is supposed to be authoritative.
Useful? React with 👍 / 👎.
| # 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" ;; |
There was a problem hiding this comment.
For a newly supported pinned gfx override such as UNSLOTH_TORCH_INDEX_URL=.../gfx1151, this only raises the torch constraint; the install/repair commands still pass bare torchvision torchaudio to uv with --index-url "$TORCH_INDEX_URL". I checked uv pip install --help: --index-url selects the package index and extra indexes require --extra-index-url, so the AMD gfx index can resolve newer independently published companions than the torch 2.11 wheel (the PS/Python paths pin these for the same ABI-mismatch reason).
Useful? React with 👍 / 👎.
| if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { | ||
| $expectedTorchTag = "rocm" |
There was a problem hiding this comment.
When a user changes an explicit ROCm pin from one ROCm family to another, for example from an existing +rocm6.4 venv to UNSLOTH_TORCH_INDEX_URL=.../gfx1151, this collapses both sides to the generic rocm flavor. If the installed package version is already current, the later dependency fast-path skips the PyTorch install entirely because the venv is not CPU-only, so the requested pinned index is never applied.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0275432536
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| case "$_torch_index_leaf" in | ||
| rocm7.2|gfx120x-all|gfx1151|gfx1150) |
There was a problem hiding this comment.
When users pin the actual AMD RDNA4 index leaf (gfx120X-all, as used by the repo's arch maps and AMD URL) via UNSLOTH_TORCH_INDEX_URL=https://repo.amd.com/rocm/whl/gfx120X-all or a mirrored family, this case-sensitive match falls through because _torch_index_leaf is not lowercased. The install then keeps the default torch>=2.4,<2.11.0 and bare companion specs, so the explicit pin can resolve the pre-2.11 wheels that this block is meant to avoid for gfx120X.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. The leaf is now lowercased before the floor case (and in the flavor/repairable helpers), so a gfx120X-all pin hits the 2.11 floor. Added bash + python tests for the capital leaf.
| return @{ | ||
| Expected = "rocm(torch<2.11)" | ||
| Installed = if ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } |
There was a problem hiding this comment.
For a pinned non-2.11 gfx index such as UNSLOTH_TORCH_INDEX_URL=https://repo.amd.com/rocm/whl/gfx110X-all, this branch treats any torch<2.11 version as a satisfied ROCm install, even if the version is an untagged CPU/CUDA wheel like 2.10.0. On a standalone studio update where the package is already current, the stale check can therefore leave $SkipPythonDeps true and never apply the explicit ROCm pin; require an installed +rocm tag before returning the expected rocm(torch<2.11) state.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. Get-RocmPinStaleTags now requires an installed +rocm tag for a non-2.11 gfx pin, so an untagged 2.10.0 wheel is reported stale and the ROCm pin is applied. Mirrored the same rule in the python _rocm_pin_family_mismatch. Added ps + python tests.
| elif _idx_leaf.startswith("cu"): | ||
| _TORCH_BACKEND = "cuda" |
There was a problem hiding this comment.
When studio update is run with a full override URL whose last segment is current, this startswith("cu") check classifies the unknown mirror as CUDA even though the new _is_cuda_family_leaf() helper intentionally rejects current/custom. That sets _TORCH_BACKEND="cuda"; on AMD hosts _ensure_rocm_torch() then returns before probing, so an existing CPU or wrong torch build is not repaired despite the override. Use the same ^cu[0-9] helper here.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. The _TORCH_BACKEND derivation now uses _is_cuda_family_leaf (^cu[0-9]), so current/custom fall through to backend "" and the GPU is probed. Also aligned install.sh to brand cuda only on cu[0-9]* (unset otherwise). Added tests.
| return @{ | ||
| Expected = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } | ||
| Installed = if ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } |
There was a problem hiding this comment.
For an explicit rocmX.Y pin, this fallback only compares the torch release line when the installed wheel has no +rocm tag. If a venv currently has an untagged CPU/CUDA wheel on the same side of the 2.11 boundary (for example UNSLOTH_TORCH_INDEX_FAMILY=rocm7.2 with torch==2.11.0), the stale check reports the pin as satisfied and the fast update path can skip dependency work, leaving the requested ROCm wheel unapplied. Treat a missing installed ROCm tag as stale for rocmX.Y pins.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. The rocmX.Y fallback now treats a missing +rocm tag as stale, so an untagged 2.11.0 wheel no longer satisfies a rocm7.2 pin. Same rule applied to the python side. Added ps + python tests.
…eaf classification Normalize torch-index leaves to lowercase before the gfx*/rocm*/cu* allowlist matches so the canonical gfx120X-all (capital X) gets the torch 2.11 floor in install.sh (leaf, flavor and repairable helpers). Require an installed +rocm local tag before a rocmX.Y or non-2.11 gfx pin is judged satisfied in setup.ps1 Get-RocmPinStaleTags and the Python _rocm_pin_family_mismatch, so an untagged CPU/CUDA wheel never leaves the pin unapplied. Classify a leaf as CUDA only via ^cu[0-9]: the Python _TORCH_BACKEND derivation now uses _is_cuda_family_leaf, and install.sh brands cuda only on cu[0-9]* (unset on an unknown /current /custom mirror leaf) so the stack probes the GPU instead of skipping ROCm repair. Add bash, Python and PowerShell tests for capital gfx120X-all floor, current/custom not-cuda, and untagged-wheel ROCm pins.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1213135cd1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if leaf in _ROCM_GFX_TORCH211_LEAVES: | ||
| # Expect the AMD per-arch wheel (three-part +rocmA.B.C, torch 2.11+). | ||
| # A satisfied per-arch install is NOT a mismatch (no reinstall loop); | ||
| # a generic rocm wheel or any pre-2.11 build IS a mismatch even at 2.11. | ||
| return not (_inst_is_211 and _inst_is_perarch) |
There was a problem hiding this comment.
When a venv was previously installed from one AMD per-arch index and the user changes the explicit pin to another 2.11 gfx index (for example gfx1151 -> gfx120X-all), this accepts any installed 2.11+ wheel with a three-part +rocmA.B.C tag as satisfying the new pin. The arch-specific wheel/dependency set is therefore not reinstalled, leaving the environment on the old GPU-family index even though the override is supposed to be authoritative.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. Added a per-venv torch-index marker (.unsloth-torch-index) written after each successful torch install with the exact resolved --index-url. On update the pin is compared against the marker by exact string match (normalize: trim, strip trailing slashes, lowercase the leaf), so a gfx1151 -> gfx120X-all switch now reinstalls even though both wheels carry the same +rocm7.13.0 tag. The wheel version tag is no longer consulted for "did the pinned index change?". No marker -> falls back to the existing heuristic (backward compatible); matching marker -> no reinstall loop.
| if url is None: | ||
| return None | ||
| leaf = url.rstrip("/").rsplit("/", 1)[-1].lower() | ||
| return url if leaf.startswith(("rocm", "gfx")) else None |
There was a problem hiding this comment.
For an explicit UNSLOTH_TORCH_INDEX_URL whose final segment is not rocm* or gfx* (for example a private mirror ending in /simple or /current), this returns None, so _ensure_rocm_torch() falls back to GPU probing and reinstalls from the auto-detected PyTorch ROCm index instead of the pinned URL. On AMD hosts where the current torch is CPU/CUDA or gets overwritten during the same install, that violates the documented “URL wins, verbatim” override and can pull wheels from the public/default index despite the pin.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. Added _ensure_verbatim_torch_index: for an explicit UNSLOTH_TORCH_INDEX_URL whose leaf is not rocm/gfx/cu/cpu (a private mirror ending in /simple, /current, ...), when the marker differs from the pin it reinstalls torch/vision/audio verbatim from that URL (exclusive --index-url, bare specs). With no marker it is a no-op so an old venv is not blindly force-reinstalled. This honors "URL wins verbatim" instead of returning None and re-probing the GPU.
| case "$_torch_index_leaf" in | ||
| rocm7.2|gfx120x-all|gfx1151|gfx1150) | ||
| 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" |
There was a problem hiding this comment.
This case only raises the torch/vision/audio constraints for rocm7.2, so an explicit pin to a newer ROCm family such as UNSLOTH_TORCH_INDEX_FAMILY=rocm7.3 or a custom /rocm8.0 URL still installs with the default torch>=2.4,<2.11.0 constraints. Those pinned ROCm indexes are treated as 2.11+ in the PowerShell/Python paths, and on the shell installer this mismatch can make the explicit override fail resolution or select the older wheel line instead of the requested ROCm family.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. Aligned the KNOWN-2.11 rocm/gfx set to exactly rocm7.2 plus the gfx allowlist (gfx120x-all, gfx1151, gfx1150) across all four installers, defined once per language: install.sh's 2.11 floor case, _ROCM_KNOWN_TORCH211_VERSIONS in install_python_stack.py, Test-RocmKnown211Version in setup.ps1, and the pinned-ROCm floor in install.ps1. No side floors an unknown newer rocm (rocm7.3 does not exist) speculatively in either the install-spec or the stale/mismatch path.
Introduce a torch-index MARKER that records the exact wheel --index-url used
after each successful torch install, so `unsloth studio update` / repair makes
the "did the pinned index change?" decision by an EXACT string compare rather
than inferring it from the wheel +rocm/+cu version tag. The tag cannot encode
the AMD per-arch gfx family (two 2.11 gfx indexes both install +rocm7.13.0), so
the tag heuristic missed a gfx1151 -> gfx120X-all switch and a custom-URL swap.
Marker path is per-venv (.unsloth-torch-index), one line = the resolved index
URL, written atomically (temp + rename). Path, format and normalization are
shared across all four installers (install.sh, install_python_stack.py,
setup.ps1, install.ps1).
- Reapply gfx pins on a per-arch target change: the marker's exact compare
reinstalls when the pinned index differs, even when both wheels share a tag.
- Honor custom ROCm URL pins during repair: an explicit index whose leaf is not
rocm/gfx/cu/cpu (e.g. simple, current) now reinstalls torch VERBATIM from the
pin when it differs from the marker ("URL wins verbatim").
- Align the KNOWN-2.11 rocm/gfx set to exactly rocm7.2 plus the gfx allowlist
gfx120x-all/gfx1151/gfx1150 in every language; stop treating an unknown newer
rocm (rocm7.3, which does not exist) as the 2.11 line speculatively.
Backward compatible: with no marker (old venvs, torch installed out-of-band) the
existing +rocm/version-tag heuristics still decide, and a matching marker never
reinstall-loops. A cu128 CUDA pin stays a CUDA pin; custom and current leaves are
not CUDA. Adds marker tests (py/sh/ps) plus cross-installer parity checks.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 35c62da452
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| $expectedTorchTag = "pinned:$_pinnedIdx" | ||
| $installedTorchTag = "pinned:$_pinnedIdx" |
There was a problem hiding this comment.
When a pin is set and the marker matches it, this branch overwrites the actual $installedTorchTag with the pinned value, so a direct unsloth studio update can miss a stale wheel. For example, if .unsloth-torch-index still records .../cu128 but torch was later replaced with a CPU wheel, the package-version fast path can set $SkipPythonDeps = $true and skip the later repair code entirely, leaving the CPU torch installed despite the CUDA pin. The marker should be an additional pin-change signal, not a substitute for validating the installed flavor for known cu*/cpu/rocm leaves.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6bbc692. The marker compare is now an additional rebuild trigger rather than a substitute: a matching marker no longer overwrites the detected installed flavor, so a stale +cpu wheel under a cuXXX pin is still caught by the flavor check and reinstalled.
| # exact string compare rather than the version-tag heuristic. Only when torch was | ||
| # actually installed from a resolved index (skip --no-torch / no-URL fallback). | ||
| if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then | ||
| _write_torch_index_marker "$VENV_DIR" "$TORCH_INDEX_URL" |
There was a problem hiding this comment.
When the Linux Radeon branch above succeeds, torch is installed from explicit Radeon wheel URLs discovered via --find-links $_RADEON_BASE_URL, not from $TORCH_INDEX_URL. Recording $TORCH_INDEX_URL therefore makes the marker claim the generic PyTorch ROCm index was used; if the user later pins that same ROCm family to move away from Radeon wheels, the marker comparison sees a match and skips the reinstall even though the installed wheels came from a different source. Record the actual Radeon source (or leave the marker absent) for that path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6bbc692. The Radeon --find-links path now records its repo.radeon.com base in the marker instead of the generic pytorch.org ROCm fallback index, so a later pin to that same ROCm family reinstalls instead of comparing equal. Mirrors install.ps1/setup.ps1, which already record the real AMD index.
| if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) { | ||
| $expectedTorchTag = "rocm" |
There was a problem hiding this comment.
On a direct Windows unsloth studio update for a supported AMD GPU where the existing venv still has CPU torch (the upgrade case handled later by the $script:ROCmGfxArch repair path), setting the expected tag to rocm marks the venv stale before that repair block runs. In the non-installer-managed path this removes $VenvDir and then immediately hits the “Virtual environment not found” error instead of reinstalling ROCm torch, so users upgrading from older CPU-only installs can lose the Studio venv and be told to rerun the full installer.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6bbc692. A supported AMD arch carrying CPU torch now expects cpu here, so the venv is reused and the ROCm override below upgrades it to ROCm in place instead of wiping it and hitting 'Virtual environment not found'. A genuinely wrong CUDA wheel still mismatches rocm and rebuilds.
Three narrow fixes in the marker-based stale-venv detection: - setup.ps1: a matching marker no longer overwrites the detected installed flavor. The marker compare is now an additional rebuild trigger, so a stale wheel (torch swapped to a +cpu build while the marker still records a cuXXX pin) is still caught by the flavor check instead of being masked as up to date. - setup.ps1: a supported AMD arch carrying CPU torch is no longer marked stale and wiped. The downstream AMD Windows ROCm override upgrades CPU torch to ROCm in place, so wiping first would delete the venv and abort with "Virtual environment not found". Only a genuinely wrong CUDA wheel still rebuilds. - install.sh: the Radeon --find-links path records its repo.radeon.com base in the marker instead of the generic pytorch.org ROCm fallback index, so a later pin to that generic family correctly reinstalls rather than comparing equal. Mirrors install.ps1/setup.ps1, which already record the real AMD index.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 257ca41753
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _ensure_cuda_torch() | ||
| _ensure_rocm_torch() | ||
| _ensure_cpu_torch() | ||
| _ensure_verbatim_torch_index() |
There was a problem hiding this comment.
When UNSLOTH_TORCH_INDEX_URL points at a private/simple index whose final segment is not cu*, cpu, rocm*, or gfx*, the fresh installer writes a matching marker, but on an AMD/ROCm host _ensure_rocm_torch() still runs first because the pin is classified as unknown. If that custom index installed a CPU/CUDA wheel (a valid cross-install/custom-mirror use case for the new verbatim URL override), _ensure_rocm_torch() can replace it with the auto-detected ROCm index; this later _ensure_verbatim_torch_index() call then sees the marker already matches the custom URL and returns without restoring it, so the full URL override is not actually honored.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 5c93ffd. _ensure_cuda_torch and _ensure_rocm_torch now return early when an explicit custom-index pin names no known torch family, so the verbatim URL override wins and is applied by _ensure_verbatim_torch_index instead of being clobbered by auto-detected CUDA/ROCm wheels first.
| _marker_verdict = _marker_pin_mismatch(_rocm_pin) | ||
| if _marker_verdict is None: | ||
| _rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver) | ||
| else: | ||
| _rocm_pin_mismatch = _marker_verdict |
There was a problem hiding this comment.
When a marker exists and equals the explicit ROCm/gfx pin, this branch sets _rocm_pin_mismatch to False without checking _installed_torch_ver. If torch was changed after the marker was written (for example the marker still records gfx1151 but the venv now contains generic 2.11.0+rocm7.2, or an older +rocm6.4 wheel), rocm_torch_ready becomes true and the pinned index is never reinstalled; setup.ps1 explicitly keeps the flavor/version check in addition to the marker for this reason.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 5c93ffd. The ROCm marker is now additive rather than a substitute: a matching marker still runs _rocm_pin_family_mismatch against the installed version, so a wheel swapped after the marker was written (marker records gfx1151 but the venv now has generic +rocm7.2 or older +rocm6.4) is caught. Mirrors setup.ps1.
| $_markerMismatch = Test-MarkerPinMismatch -VenvDir $VenvDir -PinUrl $_pinnedIdx | ||
| if ($_markerMismatch -eq $true) { $shouldRebuild = $true } |
There was a problem hiding this comment.
In a direct Windows unsloth studio update where the user changes UNSLOTH_TORCH_INDEX_URL/_FAMILY, a marker mismatch now sets $shouldRebuild, but the non-installer-managed rebuild path removes $VenvDir and then immediately fails at the later “Virtual environment not found” check instead of applying the new pin. This makes the new update-time pin override unusable unless the user reruns the full installer.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 5c93ffd. 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, so a direct unsloth studio update applies the new pin rather than hitting 'Virtual environment not found'. A broken venv or unpinned GPU-detected drift still wipes and delegates to install.ps1.
| $MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } else { $TorchIndexUrl } | ||
| Write-TorchIndexMarker -VenvDir $VenvDir -IndexUrl $MarkerIndexUrl |
There was a problem hiding this comment.
When a pinned ROCm/gfx install fails in install.ps1, the code installs a CPU fallback from $CpuFallbackIndexUrl and clears $ROCmIndexUrl, but $TorchIndexUrl still contains the original ROCm pin. This marker write therefore records the ROCm URL even though CPU torch was installed; the subsequent installer-managed setup sees a CPU wheel under a ROCm pin and aborts as stale instead of continuing with the intended CPU base/retry path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 5c93ffd. When a pinned ROCm install fails over to a CPU base, the marker now records the CPU fallback index actually used instead of the ROCm pin still in $TorchIndexUrl, so the next installer-managed setup does not see CPU torch under a ROCm pin and abort as stale.
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.
5c93ffd folded the pin-change force-reinstall into the ROCm CPU-fallback condition on one line, so the exact literal that test_pr5940_followups.py checks (if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }) no longer appeared and the test failed. Split the two conditions into separate if lines: the ROCm fallback line is restored verbatim and the pin-change force is its own line. Both still set $cpuForce to the array, so @Splat passes one arg.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0d48584e8f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # would compare-equal and skip the reinstall even though the wheels | ||
| # came from a different source. Mirrors install.ps1/setup.ps1, which | ||
| # record $ROCmIndexUrl (the real AMD index) for their AMD path. | ||
| _TORCH_MARKER_INDEX_URL="$_RADEON_BASE_URL" |
There was a problem hiding this comment.
When the Radeon --find-links install succeeds, this records _RADEON_BASE_URL, but if the subsequent unsloth install strips the HIP build the repair block below reinstalls torch from the generic $TORCH_INDEX_URL without clearing this variable. In that scenario the final .unsloth-torch-index marker still claims the venv contains Radeon wheels, so a later update/pin to the Radeon URL can compare equal and skip the reinstall even though the actual torch was repaired from the generic PyTorch ROCm index.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. The ROCm CPU repair reinstalls from the generic $TORCH_INDEX_URL, so it now sets _TORCH_MARKER_INDEX_URL to that index. The marker no longer misreports the Radeon repo after a repair, so a later Radeon pin will not compare equal and skip the reinstall.
| _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 != _pin_leaf: |
There was a problem hiding this comment.
When a standalone unsloth studio update changes UNSLOTH_TORCH_INDEX_URL from one CUDA index to another with the same leaf (for example official /cu128 to an internal mirror /cu128), this check only compares the installed +cu128 tag and then returns without consulting the .unsloth-torch-index marker. The result is that the new full URL is never used or recorded, even though this change advertises exact URL pins as authoritative.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. _ensure_cuda_torch now consults the exact-URL marker (_marker_pin_mismatch) when the installed +cuXXX tag matches the pinned leaf, so a same-leaf CUDA mirror change is reinstalled from the pinned URL (via _detect_cuda_torch_index_url) and re-recorded instead of skipped.
| head, sep, leaf = url.rpartition("/") | ||
| if sep: | ||
| return f"{head}/{leaf.lower()}" | ||
| return url.lower() |
There was a problem hiding this comment.
For custom full-URL pins whose final path segment is not a known wheel family, lowercasing the leaf makes two different case-sensitive index URLs compare equal, e.g. a marker for https://mirror.local/Current and a new pin to https://mirror.local/current. In that case _ensure_verbatim_torch_index() treats the marker as matching and skips the reinstall, even though the URL override is supposed to be used verbatim; restrict the leaf lowercasing to known family leaves like gfx120X-all/cu128 or compare custom URLs case-sensitively.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. Lowercasing is now restricted to known-family leaves (rocm/gfx/cpu/cuXXX) via _normalize_family_leaf, mirrored in install.sh and setup.ps1. A custom unknown-family leaf keeps its case, so /Current and /current no longer compare equal. Tests updated to assert this.
…rker Address three Codex review findings on the torch-index marker mechanism: - install.sh: after the ROCm CPU repair reinstalls torch from the generic $TORCH_INDEX_URL, record that as the marker source. A Radeon --find-links install set _TORCH_MARKER_INDEX_URL to its repo.radeon.com base earlier, so leaving it made the marker misreport Radeon wheels and a later Radeon pin would compare equal and skip a needed reinstall. - install_python_stack.py: _ensure_cuda_torch now consults the exact-URL marker (_marker_pin_mismatch) when the installed +cuXXX tag matches the pinned leaf, so a same-leaf CUDA mirror change (official cu128 to an internal cu128 mirror) is reinstalled and re-recorded instead of skipped. - _normalize_index_url / _normalize_family_leaf (install.sh, setup.ps1, install_python_stack.py): lowercase only KNOWN wheel-family leaves (rocm/gfx/ cpu/cuXXX) so gfx120X-all still matches gfx120x-all, while a custom (unknown-family) leaf keeps its case so a verbatim URL pin like /Current does not compare equal to /current. Tests updated to assert the refined behavior.
2 participants
Summary
get_torch_index_url(and itsstudio updatecounterpart_detect_cuda_torch_index_url) choose the PyTorch wheel family purely by probing the host GPU, with no way to override it. That is correct for a normal desktopcurl | shinstall, but it breaks any headless / container / CI install where the build host's GPU must not decide the wheel family.Concretely, in a GPU-less Docker build the builder still sees the host's NVIDIA driver through the
/proc/driver/nvidia/gpus/sysfs fallback, so_has_usable_nvidia_gpureturns true ("NVIDIA GPU detected"), but there is no queryablenvidia-smiCUDA version, so the function hits its last-ditch branch:and installs a
+cu126torch even when the image is meant to be+cu128. There was no env var to force the family.Fix
Add an explicit override, checked before any GPU probing, to both the shell installer and the Python
studio updatepath:UNSLOTH_TORCH_INDEX_URL-- a full index URL, used verbatim (highest precedence).UNSLOTH_TORCH_INDEX_FAMILY-- the convenience form (cpu,cu124,cu126,cu128,cu130,rocm6.4, ...), appended to the mirror base soUNSLOTH_PYTORCH_MIRRORis still honoured.This is the same "tell the build, don't ask the hardware" approach the published GPU images take: vLLM and SGLang both pick the CUDA wheel index from an explicit build
ARG(CUDA_VERSION->cu1xx) rather than detecting it, and the Unsloth Docker base image already pins--extra-index-url .../cu128directly. This change lets the Studio installer participate in that contract instead of guessing. Normal desktop installs are unaffected: with no override set, detection runs exactly as before.Tests
tests/sh/test_get_torch_index_url.shgains cases for: family override with no GPU, family override winning over a detected12.6(the exact Docker case), full-URL override winning over detection, family override appended toUNSLOTH_PYTORCH_MIRROR, trailing-slash stripping, URL-over-family precedence, and empty overrides falling through to detection. Full suite passes (44/44). The Python path was verified with the same matrix.