ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Installer: allow torch 2.11.x on the CUDA install path (fresh install + studio) (#6959)

Browse source 
* Studio: allow torch 2.11.x on the CUDA install path

The CUDA torch repair path (_ensure_cuda_torch) installs torch/torchvision/
torchaudio from an exclusive --index-url, so _CUDA_TORCH_PKG_SPEC decides
exactly which torch the Studio venv gets. It was capped at torch<2.11.0, so on
a cu128/cu130 host the venv resolved torch 2.10.x even though the CUDA indexes
now publish torch 2.11.0. That left the Studio venv a torch minor behind the
torch 2.11.0 Docker base image, so the CUDA dedup step would relink base libs
under a mismatched torch.

Raise the upper bound to <2.12.0 (torchvision <0.27.0, torchaudio <2.12.0) so
the CUDA install path lands on torch 2.11.x, matching the rocm7.2 spec and the
base image. The torchao selector already maps torch 2.11 -> torchao 0.17.0, and
_ensure_flash_attn degrades gracefully when no prebuilt wheel matches (Blackwell
skips it outright; non-Blackwell prints a warning and continues), so no other
pin needs to move.

Add test_cuda_torch_spec.py to lock the bound (torch 2.11.x in, 2.12.x out) and
assert the CUDA and rocm7.2 upper bounds stay in lockstep.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: use zip(strict=True) so a spec length mismatch fails loudly

* install.sh: widen the CUDA torch ceiling to <2.12.0 so a fresh install matches the base

Raising _CUDA_TORCH_PKG_SPEC alone was not enough: that spec only feeds
_ensure_cuda_torch(), the ROCm-poisoning repair path that early-returns on a
normal NVIDIA host. A fresh CUDA install (including the studio Docker build,
which runs `bash install.sh --local`) takes its torch from install.sh's
TORCH_CONSTRAINT, which was still capped at torch>=2.4,<2.11.0, so cu12x/cu13x
resolved torch 2.10.x and the venv landed a minor behind the torch 2.11.0 base
image.

Extend the existing `case "$TORCH_INDEX_URL"` block (which already relaxes
rocm7.2) with a `*/cu[0-9]*` branch that widens the ceiling to <2.12.0, keeping
the >=2.4 floor so an older CUDA index (e.g. cu118) that tops out below 2.11
still resolves. The CPU wheel and older ROCm tags stay on <2.11.0 (the glob
does not match /cpu). torchvision/torchaudio are bare on this install line and
resolve their compatible companions via wheel metadata, matching the rocm7.2
pattern.

Add behavioral tests (Python + shell) exercising the case block: cu118/124/126/
128/130 widen to <2.12.0, rocm7.2 stays 2.11.x, and /cpu plus older ROCm keep
the default <2.11.0.

* install.sh: key the CUDA torch widening off the index leaf, not the full URL

The `*/cu[0-9]*` glob matched a `cu<digit>` segment anywhere in TORCH_INDEX_URL,
so a custom UNSLOTH_PYTORCH_MIRROR whose base path contains e.g. cu128 but whose
final leaf is cpu or an older ROCm tag would still widen TORCH_CONSTRAINT to
<2.12.0, contradicting the block's own comment and letting a CPU / older-ROCm
mirror resolve torch 2.11.x. Match on _torch_index_leaf (the final path segment
the backend classification just above already computes) so only a real cu*/
rocm7.2 leaf is affected; cpu and older ROCm keep the default <2.11.0. Update
the Python + shell tests to mirror the leaf-anchored case and add regression
cases for a mirror base that contains cu128 but resolves to a cpu / rocm7.1 leaf.

* install: freeze the torch trio during the with-deps unsloth installs

Released unsloth wheels can pin an older torch than Step 1 installed
(unsloth 2026.7.2 declares torch<2.11.0), so the with-deps resolve from
PyPI silently downgrades the pinned +cuXXX torch trio to PyPI's default
wheel. The flavor guard cannot catch every such swap: PyPI's torch 2.10
default is itself cu128-flavored, so the cuXXX tag comparison still
matches while the version silently drops. Freeze the just-installed trio
with uv --overrides (overrides replace dependency requirements during
resolution), keeping torch 2.11.0+cuXXX in place while unsloth's other
dependencies resolve normally. Verified on the cu128 path: without the
override torch drops 2.11.0+cu128 -> 2.10.0; with it the trio survives
and unsloth 2026.7.2 + unsloth-zoo install cleanly.

* install: fold UV_OVERRIDE env files into the torch-trio overrides file

The CLI --overrides flag is the command-line form of UV_OVERRIDE, so
passing it replaced any overrides file already exported for the process;
macOS arm64 exports UV_OVERRIDE=overrides-darwin-arm64.txt for the same
generic install path and would have lost those pins. Concatenate any
UV_OVERRIDE files into the temp trio file so both keep applying.

* install: extend the torch-trio overrides guard to migrated installs

Four follow-ups to the Step-2 --overrides guard, all empirically verified:

1. The migrated-environment with-deps unsloth install resolved
   unsloth>=2026.7.2 (which pins torch<2.11.0) without the overrides file,
   so a migrated CUDA venv on torch 2.11 was silently downgraded -- the
   exact bug this branch fixes on the fresh path. The overrides build is
   now a function (_build_unsloth_torch_overrides, reading the trio
   installed at call time) invoked by both with-deps paths; the migrated
   no-torch path installs --no-deps and stays unguarded.

2. The overrides temp file is now cleaned by the EXIT trap (same pattern
   as _UV_OVERRIDE_TMPDIR, pre-initialized empty so an inherited value can
   never reach the trap's rm); previously any Step-2 failure leaked it.

3. Folding UV_OVERRIDE files used cat, which joins the last requirement of
   a file lacking a trailing newline onto the next file's first requirement
   (reproduced: idna==3.10certifi==2025.1.31 makes uv fail parsing).

4. Inherited torch/torchvision/torchaudio override lines are now filtered
   out when folding: uv intersects duplicate overrides rather than
   last-wins (verified on uv 0.10.12: direct conflict is unsatisfiable,
   transitive conflict silently backtracks), so a conflicting inherited
   trio pin would break the resolve the generated exact pins protect.
   Both 3 and 4 are handled by a single newline-terminating awk filter
   that preserves non-trio overrides (torchmetrics, torchao, ...).

test_unsloth_torch_override.sh extended: migrated-path coverage, trap
assertion, and a functional fold test (14 checks).

* installer: tighten comments

* install: keep the existing torch release when re-running the installer

Re-running `curl -fsSL https://unsloth.ai/install.sh | sh` over an existing
install rebuilds the venv for clean state, which silently moved users to the
newest torch in range (2.10 -> 2.11 once the constraint widened). A torch the
user already validated must survive an unsloth update.

Before the old venv is moved aside for rollback, its torch version is probed
(last stdout line only, so sitecustomize noise cannot corrupt it). After the
index leaf is chosen, _previous_torch_pin turns that version into a
torch==X.Y.Z pin, but only when it cannot do harm:

- cu*/cpu leaves only; rocm leaves keep their floors (rocm7.2 must land 2.11
  for the Strix _grouped_mm fix) and the Radeon wheel-matching path is
  untouched.
- The wheel's flavor tag must match the freshly chosen leaf, so a flavor
  change (cpu -> cuda, cu126 -> cu130) still installs the correct new build.
- The base must look like a release, so probe noise never becomes a pin.
- UNSLOTH_TORCH_UPGRADE=1 opts out and restores the old always-newest
  behavior; the substep line advertises it.

The supported range is kept in _PREV_FALLBACK_CONSTRAINT: if the exact
release is not resolvable from the chosen index (custom mirrors prune old
wheels), the install warns and falls back to the newest supported release
instead of failing the whole run. The later flavor-mismatch repair reuses
TORCH_CONSTRAINT, so a mid-install clobber is repaired back to the kept
release rather than the newest one.

Verified end to end: a venv seeded with torch 2.10.0+cu130 re-run through the
full installer finishes with torch 2.10.0+cu130 (previously 2.11.0+cu130).

Tests: tests/sh/test_previous_torch_pin.sh covers keep/flavor-change/rocm/
noise/opt-out plus wiring (probe ordering before venv replacement, fallback
present, SKIP_TORCH gate).

* install: constrain kept torch pins to the supported window

Review caught that _previous_torch_pin pinned the previous venv's torch on
flavor match alone, so a release outside the installer's active range (a
2.3.x manual install below the >=2.4 floor, or a 2.12.x manual upgrade above
the ceiling) replaced the bounds computed just above it and a rerun kept a
torch the installer otherwise deliberately excludes.

New _torch_release_in_window checks the probed base against the active
TORCH_CONSTRAINT ("torch>=A.B[,<C.D.F]") at major.minor granularity, which
is exact for the windows this script uses (ceilings are always X.Y.0; a
non-.0 ceiling would only make it conservative). Anything unparseable
answers no, so probe noise or a malformed window fails toward the supported
range instead of becoming a pin. _previous_torch_pin takes the active
constraint as a third argument and refuses out-of-window releases; the
in-window keep behavior is unchanged.

Tests: out-of-window rows (2.3.x floor, 2.12.x ceiling, boundary keeps, cpu
and macOS windows, malformed/empty windows) plus direct
_torch_release_in_window coverage.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-19 06:19:29 -07:00 • committed by GitHub
commit a9be36830e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 558 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

80
tests/python/test_tokenizers_and_torch_constraint.py
View file

@ -69,6 +69,21 @@ class TestStructuralTorchConstraint:
def test_tightened_assignment_exists(self):
assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh
def test_cuda_constraint_widened_to_2_12(self):
"""A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x
land torch 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC);
without it cu128/cu130 resolves torch 2.10.x."""
assert 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' in self._sh
def test_cuda_case_widens_via_index_leaf(self):
"""The cu* branch of the _torch_index_leaf case sets the widened
constraint (parallel to rocm7.2), anchored on the leaf."""
m = re.search(
r'cu\[0-9\]\*\)\s*TORCH_CONSTRAINT="torch>=2\.4,<2\.12\.0"',
self._sh,
)
assert m is not None, "CUDA (cu*) TORCH_CONSTRAINT widening case not found"
def test_variable_used_in_pip_install(self):
"""$TORCH_CONSTRAINT must appear in a uv pip install line."""
assert '"$TORCH_CONSTRAINT"' in self._sh
@ -384,6 +399,71 @@ class TestTorchConstraintShell:
logged = log_file.read_text()
assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}"
# Mirrors the _torch_index_leaf case in install.sh: rocm7.2 -> 2.11.x floor,
# CUDA -> widened <2.12.0 ceiling, else (CPU/older ROCm) -> default. Anchored
# on the final path segment, so a mirror base path containing cu*/rocm7.2 but
# ending in a cpu/older-rocm leaf keeps the default.
_INDEX_SNIPPET = textwrap.dedent(r"""
#!/bin/bash
set -e
TORCH_INDEX_URL="{index_url}"
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
_torch_index_leaf="${TORCH_INDEX_URL%/}"
_torch_index_leaf="${_torch_index_leaf##*/}"
case "$_torch_index_leaf" in
rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
esac
echo "$TORCH_CONSTRAINT"
""").strip()
def _resolve_index(self, tmp_path: pathlib.Path, index_url: str) -> str:
script_file = tmp_path / "index_snippet.sh"
script_file.write_text(self._INDEX_SNIPPET.replace("{index_url}", index_url))
script_file.chmod(0o755)
result = subprocess.run(
["bash", str(script_file)],
capture_output = True,
text = True,
timeout = 10,
)
assert result.returncode == 0, f"Script failed: {result.stderr}"
return result.stdout.strip()
@pytest.mark.parametrize("leaf", ["cu118", "cu124", "cu126", "cu128", "cu130"])
def test_cuda_index_widens_to_2_12(self, tmp_path, leaf):
url = f"https://download.pytorch.org/whl/{leaf}"
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0"
def test_rocm72_index_uses_211_floor(self, tmp_path):
url = "https://download.pytorch.org/whl/rocm7.2"
assert self._resolve_index(tmp_path, url) == "torch>=2.11.0,<2.12.0"
def test_cpu_index_keeps_default(self, tmp_path):
# /cpu must NOT match the */cu[0-9]* branch.
url = "https://download.pytorch.org/whl/cpu"
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
def test_older_rocm_index_keeps_default(self, tmp_path):
url = "https://download.pytorch.org/whl/rocm7.1"
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
def test_cuda_index_custom_mirror_widens(self, tmp_path):
url = "https://internal.example.com/pytorch/cu128"
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0"
@pytest.mark.parametrize(
"url",
[
"https://internal.example.com/pytorch/cu128/cpu",
"https://internal.example.com/cu128/whl/rocm7.1",
],
)
def test_cuda_in_mirror_path_but_noncuda_leaf_keeps_default(self, tmp_path, url):
# A cu128 in the mirror base path must not widen when the leaf is cpu /
# older ROCm: the case anchors on _torch_index_leaf, not the whole URL.
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
# Group 3 -- E2E tokenizers fix (requires network, ~2-5 min)
@pytest.mark.e2e