Compare commits

...
Sign in to create a new pull request.

18 commits

Author SHA1 Message Date
Daniel Han
90fc453e8b install.ps1: remove the torch overrides temp file from the outer finally
New-UnslothTorchOverridesFile writes a uv --overrides file that also folds in the
caller's inherited UV_OVERRIDE lines, which can carry authenticated direct URLs.
The in-flow Remove-Item calls only run on the normal path, so a terminating error
or Ctrl-C between the GetTempFileName and those removals left the file in %TEMP%.

install.sh already removes its twin (_UNSLOTH_TORCH_OVERRIDES) from
_cleanup_install_temporaries on both the EXIT and the signal traps, and empties the
variable beforehand so an inherited value never reaches the rm. Mirror both halves
here: reset $script:TorchOverridesFile to $null ahead of the outer try (under
irm | iex the script scope is the caller's session) and sweep it in the finally.
The block writes nothing to the success or error stream, so a Ctrl-C cannot
truncate it.

tests/studio/test_unsloth_torch_override.ps1 is the Windows twin of
tests/sh/test_unsloth_torch_override.sh: it asserts the guarded with-deps installs,
the removal sites and the null reset, and behaviourally runs the extracted finally
body against a live temp file holding a credential-bearing override line.
Before this change 17 passed, 5 failed; after, 22 passed, 0 failed.
2026-07-27 15:57:33 +00:00
Daniel Han
b5a508c74e Merge origin/main into windows-torch-211
Resolve tests/sh/test_torch_constraint.sh by keeping main's invariant
based structure from #7503 and restating it against this branch's
composed ceiling variables.

Both sides rewrote the same assertions in the same spirit. #7503
replaced brittle occurrence counts with invariants after #7354's gfx906
branch pushed three counts up by one and left Backend CI red. This
branch had independently rewritten them because the default torch range
is no longer a literal, it is composed from _TORCH_CEILING="2.12.0".

The resolution keeps both:

- The column 0 anchor for the default assignment now reads the composed
  form TORCH_CONSTRAINT="torch>=2.4,<${_TORCH_CEILING}", so an indented
  per branch pin still cannot satisfy it. _TORCH_CEILING="2.12.0" is
  asserted separately so the anchor keeps its teeth.
- The tightened macOS py3.13 assignment is an existence check rather
  than a count, per #7503.
- Both stray literal checks are kept. They are complementary: #7503's
  sees through shell escaped quotes on an install or echoed command
  line, this branch's catches the two default ranges copied onto any
  assignment that is not TORCH_CONSTRAINT=.
- The gfx906 sub 2.11 cap assertion is kept, converted from a count to
  an indented existence check so a second hardware branch adding its own
  sub 2.11 cap is not a test edit.
- #7503's every assignment upper bound checks for torchvision and
  torchaudio are kept, widened to accept the composed ceiling alongside
  a literal bound.

install.sh, install.ps1, studio/setup.ps1 and studio/install_python_stack.py
are untouched by this merge. The gfx906 literal cap and the three ROCm
2.11 floors stay literal, and the test permits them without forcing them
to track the ceiling.

tests/sh/test_torch_constraint.sh: 49 pass, 0 fail.
2026-07-27 14:27:54 +00:00
Daniel Han
ac8bc5f9a3 Merge remote-tracking branch 'origin/main' into windows-torch-211
Three files conflicted, all on the gfx906 (MI50 / Radeon VII) work that landed
on main next to this branch's comment pass and torch 2.11 default rollout.

install.sh: kept main's _is_gfx906_bnb_skip / _gfx906_bnb_snapshot / _prune
helpers and the _gfx906_bnb_snapshot call in the migrated path, with this
branch's condensed comments. The gfx906 reroute keeps its literal
TORCH_CONSTRAINT="torch>=2.4,<2.11.0": the rocm6.3 index it routes to tops out
at torch 2.9.x, so it is now a deliberate narrower cap rather than a restatement
of the default, and the comment says so.

studio/install_python_stack.py: kept main's _runtime_is_gfx906 detection, the
rocm6.3 legacy override and the bnb skip/prune branch, again with the condensed
comments. _ROCM_TORCH_PKG_SPECS["_default"] stays at <2.11.0, which is what the
gfx906 path installs from.

tests: the two sides asserted opposite things about hard-coded torch ranges in
install.sh. This branch composes the default from _TORCH_CEILING and asserted
zero literals; main asserted the literal exists and only on a TORCH_CONSTRAINT=
assignment. Merged into the rule both wanted: a literal range is allowed only on
a curated per-index TORCH_CONSTRAINT= override, never on a pip/uv install line,
plus a separate check that the gfx906 reroute keeps its sub-2.11 cap. The same
resolution is applied to tests/sh/test_torch_constraint.sh, which asserted the
zero-literal count and would otherwise have gone red on main's new assignment.
2026-07-27 13:24:03 +00:00
Daniel Han
1e3b1c97ad install_python_stack.py: keep the ROCm 7.1 repair on the 2.11 line
The rocm7.1 leaf fell through to _ROCM_TORCH_PKG_SPECS["_default"], which caps
the trio below 2.11. Now that install.sh leaves a rocm7.1 leaf on the widened
default range, a fresh install resolves torch 2.11.0+rocm7.1 while the later
dependency pass force-reinstalled 2.10.0+rocm7.1 over it, so any repair or
`studio update` silently downgraded the environment.

download.pytorch.org/whl/rocm7.1 serves a paired 2.11 trio, verified with
uv pip compile --no-deps against that index:

  install.sh default range -> torch 2.11.0+rocm7.1, torchvision 0.26.0+rocm7.1,
                              torchaudio 2.11.0+rocm7.1
  _default repair spec     -> torch 2.10.0+rocm7.1, torchvision 0.25.0+rocm7.1,
                              torchaudio 2.10.0+rocm7.1

Give rocm7.1 its own entry carrying install.sh's default range rather than the
rocm7.2 tuple: only the _grouped_mm arches take the hard 2.11 floor, so
_ROCM_KNOWN_TORCH211_VERSIONS stays {(7, 2)}. _default keeps its literal <2.11
ceiling because rocm7.0 and older genuinely top out below it (rocm7.0 at 2.10.0,
rocm6.4 and rocm6.3 at 2.9.1, rocm6.2 at 2.5.1), and the stale index comments
are corrected to match what those indexes serve today.
2026-07-26 15:50:14 +00:00
Daniel Han
1d1bd91326 Merge remote-tracking branch 'origin/main' into windows-torch-211 2026-07-26 15:13:17 +00:00
Daniel Han
57c2a3d400 Merge remote-tracking branch 'origin/windows-torch-211' into windows-torch-211 2026-07-26 11:25:28 +00:00
Daniel Han
27c81d493b install.ps1: make the release preservation robust across sessions and failures
Three review follow-ups on the Windows port:

- The previous-release probe reads dist metadata rather than importing
  torch, matching install.sh. A broken CUDA/ROCm DLL (or an import slow
  enough to hit the timeout) made the probe return nothing, so the pin
  was never created and the installer replaced the venv with the newest
  supported release even though UNSLOTH_TORCH_UPGRADE was unset.

- The preservation state is reset at the top of the install instead of
  inside the existing-venv branch. Under the documented irm | iex flow
  the script scope IS the caller's session, so a second invocation
  against a different UNSLOTH_STUDIO_HOME with no existing venv skipped
  the branch and inherited the earlier run's release, pinning a fresh
  install to an unrelated venv's torch. The pin is reset with it, so the
  no-index fallback path cannot see a stale pin either.

- UNSLOTH_KEPT_TORCH is cleared from an outer finally around the
  installer entry point. The in-flow clears and Exit-InstallFailure
  cover the handled paths, but a terminating exception between the
  export and the end of setup left the handoff set in a session that
  outlives the installer, where a later studio setup or update would
  consume the abandoned exact pin.

Verified the metadata probe end to end (returns the installed release
through the bounded runner; a missing interpreter still yields null).
Three regression guards added to tests/studio/test_previous_torch_pin.ps1
for the metadata read, the pre-branch reset ordering and the outer
finally. Full sh, ps1 and pytest installer batteries pass (host-defaults
and the tokenizers negative-control are the known pre-existing
failures).
2026-07-26 11:25:28 +00:00
danielhanchen
96d8b99e65 Merge remote-tracking branch 'origin/main' into r7256 2026-07-25 05:14:38 +00:00
danielhanchen
a14cc540ae Merge remote-tracking branch 'origin/main' into windows-torch-211
Reconcile PR #7256 (Windows torch 2.11 + release preservation, installer
comment reduction, torch 2.11 default line) with 77 commits of main install
rewrites. Base is main's newer install semantics; this branch's still-novel
contributions are layered on top.

Key decisions:
- install.ps1: kept main's rollback lifecycle (try/finally Restore-StudioVenvRollback,
  #7342) and the installed-version report (#7265); layered this branch's torch-2.11
  allowance (torch<2.12.0 on the Windows CUDA fresh-install, CPU fallback and flavor
  repairs), the release-preservation port (Get-InstalledTorchVersionRaw / kept-release
  installs / UNSLOTH_KEPT_TORCH handoff / torch-overrides freeze) and the Exit-InstallFailure
  UNSLOTH_KEPT_TORCH clear.
- install.sh: took main's newer AMD/Strix routing (per-arch index reroute #7264/#7300,
  runtime-less gfx inference #7305, KFD detection fix #7314, Radeon 8065S regex #7290,
  signal-restore trap #7342) and the #7365 unsloth pin bump; kept this branch's
  _TORCH_CEILING/_TORCHVISION_CEILING/_TORCHAUDIO_CEILING refactor widening the default
  ceiling to torch<2.12.0.
- studio/install_python_stack.py: deferred all five code conflicts to main's Strix
  inference logic (result is AST-identical to main; only comment reductions remain).
- studio/setup.ps1: took main's comment covering the whisper.cpp dictation markers
  (#7095); the UNSLOTH_KEPT_TORCH consumption handoff auto-merged intact.

Dropped as superseded by main: this branch's stale unsloth>=2026.7.4 pins (main #7365),
the older rocm7.1->rocm7.2 Strix reroute (main #7264/#7300), the narrow Radeon 80[0-9]0S
regex (main #7290), and the pre-inference has_hip_torch gate (main's rocm_torch_ready gate,
Codex P1 #7305). The two review-item fixes (grep -E in test_torch_constraint.sh,
UNSLOTH_KEPT_TORCH clear in Exit-InstallFailure) survive.
2026-07-24 11:24:55 +00:00
danielhanchen
c845e72692 Harden torch-constraint test grep portability and clear kept-torch on failure
- test_torch_constraint.sh: use grep -E for the hardcoded-range guard. The BRE
  \| alternation is a GNU extension; on BSD/macOS grep it can be a literal, so
  the regression guard could silently no-op on a supported platform.
- install.ps1: clear UNSLOTH_KEPT_TORCH in Exit-InstallFailure. A non-Tauri
  irm | iex run throws while the caller session stays alive, so a leaked
  kept-torch handoff could let a later setup/update re-pin an abandoned exact
  torch release.
2026-07-24 11:09:12 +00:00
Daniel Han
0a32590f71 Merge remote-tracking branch 'origin/main' into windows-torch-211
# Conflicts:
#	install.ps1
#	install.sh
2026-07-21 02:20:56 +00:00
Daniel Han
8d3f10145c Merge remote-tracking branch 'origin/windows-torch-211' into windows-torch-211 2026-07-20 12:07:17 +00:00
Daniel Han
d90fd8b563 install: harden the preservation probe and the Windows override handoff
Three review follow-ups:

- install.sh's existing-venv torch probe now reads the version from dist
  metadata via the bounded runner instead of importing torch: a wedged
  CUDA/ROCm driver can hang the import indefinitely, while the metadata
  read never touches the driver. A failed or timed-out probe simply
  yields no preservable version.

- install.ps1's New-UnslothTorchOverridesFile folds caller-supplied
  UV_OVERRIDE files into the temporary overrides file (minus their
  torch-trio lines), matching install.sh: --overrides replaces the env
  file wholesale, so without the merge a caller's own dependency
  overrides were silently dropped during the unsloth resolution.

- install.ps1 clears an inherited UNSLOTH_KEPT_TORCH before the pin
  decision and sets it only on a fresh Get-PreviousTorchPin result: an
  interrupted earlier run could leak a stale exact pin into setup.ps1
  even when the current run found nothing to preserve or the upgrade
  opt-in was set.

Verified: the metadata probe returns the installed release through the
bounded runner; the override filter keeps unrelated pins (numpy,
torchao) while dropping torch/torchvision/torchaudio lines in all spec
forms. Full sh, ps1 and pytest installer batteries pass (host-defaults
and the tokenizers negative-control are the known pre-existing
failures).
2026-07-20 12:07:08 +00:00
pre-commit-ci[bot]
49e304221a [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-20 10:46:17 +00:00
Daniel Han
3f79b5e53d install: make torch 2.11 the default supported line on wheel-backed platforms
Fresh installs now resolve the torch 2.11 trio (torch 2.11.0,
torchvision 0.26.0, torchaudio 2.11.0) everywhere the wheels exist,
verified by live index resolution for py3.11-3.13: Linux x86_64
(cpu, cu126, cu128, cu130, rocm7.1, rocm7.2), Linux aarch64 (cpu,
cu130), Windows x86_64 (cpu, cu126, cu128, cu130; triton-windows<3.7
resolves 3.6.0.post26, the 2.11 pairing) and macOS arm64. Existing
installs are unaffected: the release preservation on both installers
keeps the installed torch on re-runs (verified end to end against the
live cpu index: a seeded 2.10.0 venv stays 2.10.0 with paired 0.25/2.10
companions while a fresh venv resolves the 2.11 trio).

- install.sh: the supported range is centralized in _TORCH_CEILING /
  _TORCHVISION_CEILING / _TORCHAUDIO_CEILING and composed into the
  default constraints, so the next bump (torch 2.12) is a three-line
  change. The now-redundant cu* widen arm and custom-pin companion
  block collapse into the default; the curated ROCm >=2.11 floors stay
  literal.
- install.ps1: the hoisted default trio, the CUDA flavor repair and the
  ROCm CPU fallback all use the <2.12 trio (the leaf gate collapsed
  since cu and non-cu now share the range).
- setup.ps1: the unknown-leaf pinned CUDA trio widens to <2.12,
  matching the pinned cpu path.

Deliberately NOT widened: rocm6.4 (tops at torch 2.9.1) and rocm7.0
(2.10.0) by index content; macOS x86_64 (no >=2.4 wheels, stays
no-torch); the AMD per-arch repo.amd.com curated bounds for arches
outside the 2.11 allowlist.

Constraint suites updated to the ceiling scheme; parity suites updated
to the widened trio. All sh, ps1 and pytest installer suites pass (the
host-defaults suite and the tokenizers negative-control are known
pre-existing failures).
2026-07-20 10:44:42 +00:00
Daniel Han
e49a89f4f6 install.ps1: port the torch release preservation from install.sh
The Windows installer recreates the venv on every re-run (the existing
environment is moved aside for rollback and a fresh venv is created), so
irm install.ps1 | iex over an existing install resolved the newest torch
in range instead of keeping the installed release. This ports the
install.sh preservation merged in PR 7250:

- The installed torch version is probed (bounded process, drained
  streams, 30s timeout, last non-empty stdout line) BEFORE the rollback
  move, while the old interpreter still exists.
- Get-PreviousTorchPin mirrors _previous_torch_pin: numeric-strict base
  (nightly/dev/rc builds never pin), UNSLOTH_TORCH_UPGRADE=1 opt-out,
  and a release-in-window check against the route's final constraint,
  so a raised ROCm floor correctly rejects keeping an older release.
  The decision runs after every index and floor decision.
- The kept release installs as an exact pin with minor-paired
  companions (torchvision 0.minor+15, torchaudio 2.minor) at the main
  install, the AMD ROCm install and both flavor repairs; when the pinned
  release is not installable from the selected index the installer warns,
  clears the pin and falls back to the supported range. The CPU fallback
  cannot be reached with a live pin (the ROCm site resolves or clears it
  first) and keeps the plain range install.
- The with-deps unsloth installs now carry a uv overrides file with the
  exact installed trio (twin of _build_unsloth_torch_overrides), so
  dependency resolution cannot move torch after it was deliberately
  selected.
- The kept release is exported as UNSLOTH_KEPT_TORCH for setup.ps1,
  whose ROCm/CPU/CUDA torch installs substitute the kept trio when the
  env var is present (ROCm floors dominate; behavior without the env
  var is unchanged, so direct studio update runs are unaffected). After
  setup returns the installer re-probes and warns loudly if the kept
  release series changed, then clears the env var.

Preservation is version-agnostic (numeric parse + window comparison),
so a future ceiling bump to torch 2.12 keeps 2.11 installs in place.

New suite tests/studio/test_previous_torch_pin.ps1 (AST-extracted
helpers, 50+ checks incl. future-ceiling cases and structural wiring)
passes, along with the flavor/pin-hardening/pin-stale/node ps1 suites,
the parity and install-stack pytest suites, and the sh preservation and
constraint suites.
2026-07-20 10:41:41 +00:00
Daniel Han
d5b743f13e install: reduce comment volume across the installers
Collapse the multi-line rationale blocks that accreted across the torch
index, override, redaction and companion-pin work to single-line
constraints, and drop review-history narration. Comments and blank lines
only; the PowerShell files verify code-identical by token-stream
comparison against the previous commit, install.sh passes sh -n, and the
Python file passes the AST comments-only gate.

install.sh -258 lines, install.ps1 -235, studio/setup.ps1 -304,
studio/install_python_stack.py -182. One test touched:
tests/sh/test_mac_intel_compat.sh anchors its awk extraction on a
comment phrase that was shortened; the anchor now matches both wordings.

Full battery: parity/install-stack/pr5940 pytest suites, all sh suites
(host-defaults is the known pre-existing failure), and the four studio
ps1 suites all pass.
2026-07-20 10:09:37 +00:00
Daniel Han
25f6b8d842 install.ps1: allow torch 2.11 on Windows CUDA installs
The Windows installer still capped torch at <2.11.0 everywhere, a
stability pin from before 2.11 support landed. Every other layer already
allows 2.11: install.sh widens cu<digits> leaves to <2.12.0, the Python
repair layer uses the <2.12.0 trio for CUDA and CPU, and setup.ps1's
bare specs on known cu leaves resolve 2.11 today.

The fresh-install and CUDA flavor-repair paths now widen the trio to
torch<2.12 / torchvision<0.27 / torchaudio<2.12 when the index leaf is a
cu<digits> family, and keep the 2.10 line otherwise (custom pins and the
CPU fallback are unchanged, matching install.sh's defaults). Verified by
uv dry-runs against cu126/cu128/cu130 for win_amd64: the trio resolves
paired at 2.11.0 / 0.26.0 / 2.11.0, and the existing triton-windows<3.7
constraint resolves 3.6.0.post26, the torch 2.11 pairing.

Parity test updated to assert the leaf-gated widen.
2026-07-20 08:05:55 +00:00
11 changed files with 1498 additions and 2058 deletions

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -44,15 +44,10 @@ IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64"
IS_MAC_ARM = IS_MACOS and platform.machine() == "arm64"
IS_LINUX = sys.platform.startswith("linux")
# amd-smi auto-elevates on Windows (UAC/DiskPart prompt mid-install). This installer
# only spawns probes and pip/uv (no elevation), so set __COMPAT_LAYER=RunAsInvoker
# process-wide; amd-smi then runs un-elevated. setup.ps1 keeps per-call guards (it
# also spawns winget installers that need elevation).
# amd-smi auto-elevates on Windows (UAC/DiskPart prompt); RunAsInvoker keeps this installer's probes un-elevated.
if IS_WINDOWS:
os.environ.setdefault("__COMPAT_LAYER", "RunAsInvoker")
# torchcodec ships wheels only for manylinux_2_28_x86_64, macosx_12_0_arm64,
# and win_amd64. On other hosts the audio extras must be filtered out (the
# extras-no-deps step would otherwise fail), regardless of NO_TORCH.
# torchcodec ships wheels only for manylinux_2_28_x86_64, macosx_12_0_arm64, win_amd64; elsewhere filter the audio extras regardless of NO_TORCH.
PLATFORM_LACKS_TORCHCODEC_WHEEL = (
(IS_LINUX and platform.machine() in {"aarch64", "arm64"})
or (IS_WINDOWS and platform.machine().lower() in {"arm64", "aarch64"})
@ -60,12 +55,11 @@ PLATFORM_LACKS_TORCHCODEC_WHEEL = (
)
# ── ROCm / AMD GPU support ─────────────────────────────────────────────────────
# Detected ROCm (major, minor) -> best PyTorch wheel tag on
# download.pytorch.org. Checked newest-first (>=).
# Detected ROCm (major, minor) -> best PyTorch wheel tag, checked newest-first (>=).
_ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
(7, 2): "rocm7.2", # torch 2.11.0
(7, 1): "rocm7.1", # torch 2.10.0
(7, 0): "rocm7.0",
(7, 1): "rocm7.1", # torch 2.11.0
(7, 0): "rocm7.0", # torch 2.10.0
(6, 4): "rocm6.4",
(6, 3): "rocm6.3",
(6, 2): "rocm6.2",
@ -134,27 +128,36 @@ _ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset(
{"gfx120x-all", "gfx1151", "gfx1150", "gfx1152"}
)
# pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (rocm7.2 only today); don't
# floor an unknown newer rocm speculatively. Match install.sh / setup.ps1 / install.ps1.
# rocmX.Y indexes KNOWN to ship torch 2.11; never floor an unknown newer rocm speculatively.
_ROCM_KNOWN_TORCH211_VERSIONS: frozenset[tuple[int, int]] = frozenset({(7, 2)})
# Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x).
# Per-tag pip specs for the repair/update path; must land on the same wheels a fresh
# install.sh run would pick, otherwise `studio update` silently downgrades the venv.
_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
# Floored at 2.11 (the _grouped_mm bug), matching install.sh's rocm7.2|gfx* case.
"rocm7.2": (
"torch>=2.11.0,<2.12.0",
"torchvision>=0.26.0,<0.27.0",
"torchaudio>=2.11.0,<2.12.0",
),
# rocm7.1 and earlier: torch 2.x below 2.11
# rocm7.1 also serves a paired 2.11 trio (torch 2.11.0 / torchvision 0.26.0 /
# torchaudio 2.11.0), so it takes install.sh's widened DEFAULT range rather than
# the 2.11 floor: no _grouped_mm floor applies here, but capping at <2.11 would
# force-reinstall 2.10 over the 2.11 a fresh install just resolved.
"rocm7.1": (
"torch>=2.4,<2.12.0",
"torchvision>=0.19,<0.27.0",
"torchaudio>=2.4,<2.12.0",
),
# rocm7.0 and earlier genuinely top out below 2.11 (rocm7.0: torch 2.10.0,
# rocm6.4/6.3: 2.9.1, rocm6.2: 2.5.1), so the old ceiling stays literal.
"_default": (
"torch>=2.4,<2.11.0",
"torchvision>=0.19,<0.26.0",
"torchaudio>=2.4,<2.11.0",
),
}
# Windows AMD per-arch companion pins for the repo.amd.com index (mirrors the install.ps1 /
# setup.ps1 floor maps): pinning stops the per-arch index (each published independently) from
# resolving an ABI-mismatched companion. Unlisted arches have no floor, so stay bare.
# Windows AMD per-arch companion pins for repo.amd.com: pinning stops the per-arch index resolving an ABI-mismatched companion; unlisted arches stay bare.
_WINDOWS_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
"gfx1201": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
"gfx1200": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
@ -224,31 +227,18 @@ def _torch_index_leaf(url: str) -> str:
return path.rstrip("/").rsplit("/", 1)[-1].lower()
# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed (torchao
# 0.17 cpp loads cleanly, and the flash-attn/causal-conv1d/mamba wheels pass on 2.11).
# torchvision/torchaudio are pinned (not bare) so the exclusive --index-url can't
# resolve one built against a different torch major -> ABI mismatch.
# CUDA torch repair specs (see _ensure_cuda_torch): companions pinned so the exclusive --index-url can't resolve an ABI-mismatched torch major.
_CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = (
"torch>=2.4,<2.12.0",
"torchvision>=0.19,<0.27.0",
"torchaudio>=2.4,<2.12.0",
)
# CPU torch repair specs (see _ensure_cpu_torch). Same bounds/reasoning as CUDA: the
# /cpu index also serves newer torch, so a bare trio could resolve out of range or ABI-
# mismatched.
# CPU torch repair specs (see _ensure_cpu_torch): the /cpu index also serves newer torch, so a bare trio could resolve ABI-mismatched.
_CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC
# torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch
# mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails
# to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a
# range, so match torchao to the installed torch (table: pytorch/ao#2919):
# 2.9.x -> 0.14.0
# 2.10.x, CUDA<=12 -> 0.16.0 (cpp built for 2.10, loads via the CUDA-12 wheel)
# 2.10.x, CUDA>=13 -> 0.17.0 (cu130: 0.16.0's CUDA-12 cpp crashes on load; 0.17.0
# targets torch 2.11 so its cpp is cleanly skipped, not crashed)
# 2.11.x -> 0.17.0 (reachable via CUDA or ROCm rocm7.2)
# Unknown/older torch keeps the conservative default.
# torchao's cpp extensions are pinned to ONE torch release AND CUDA major (table: pytorch/ao#2919):
# 2.9.x -> 0.14.0; 2.10.x CUDA<=12 -> 0.16.0; 2.10.x CUDA>=13 -> 0.17.0; 2.11.x -> 0.17.0; else default.
_TORCHAO_DEFAULT_SPEC = "torchao==0.14.0"
_TORCHAO_TORCH_210_SPEC = "torchao==0.16.0"
_TORCHAO_TORCH_210_CUDA13_SPEC = "torchao==0.17.0"
@ -279,8 +269,7 @@ def _select_torchao_spec(torch_version: str | None) -> str:
release = str(torch_version).split("+", 1)[0] # drop +cu130/+rocm6.4/+cpu
parts = release.split(".")
try:
# Strip any pre-release/dev suffix from the minor (e.g. '10rc1' -> '10'),
# matching wheel_utils.probe_torch_wheel_env.
# Strip any pre-release/dev suffix from the minor (e.g. '10rc1' -> '10').
minor_str = re.sub(r"[^0-9].*", "", parts[1]) if len(parts) > 1 else ""
major, minor = int(parts[0]), int(minor_str)
except (IndexError, ValueError):
@ -358,9 +347,7 @@ def _installed_torch_is_windows_rocm() -> bool:
return probe.returncode == 0 and bool(lines and lines[-1] == "yes")
# constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install
# from before the cap existed can already be stuck at 4.14+, which later
# constrained installs won't touch since it already satisfies mcp/fastmcp.
# constraints.txt caps anyio <4.14 (#6483), but a pre-cap install can be stuck at 4.14+ which constrained installs won't touch.
_ANYIO_BAD_FLOOR = (4, 14)
@ -393,14 +380,12 @@ def _repair_bad_anyio() -> None:
)
# AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/).
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs.
# AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/); override with UNSLOTH_ROCM_WINDOWS_MIRROR.
_ROCM_WINDOWS_INDEX_BASE = (
os.environ.get("UNSLOTH_ROCM_WINDOWS_MIRROR") or "https://repo.amd.com/rocm/whl"
).rstrip("/")
# gfx arch → AMD index arch-family suffix; each family is a separate
# pip index on repo.amd.com.
# gfx arch → AMD index arch-family suffix; each family is a separate pip index on repo.amd.com.
_GFX_TO_AMD_INDEX_ARCH: dict[str, str] = {
"gfx1201": "gfx120X-all",
"gfx1200": "gfx120X-all", # RDNA 4
@ -422,9 +407,7 @@ _GFX_TO_AMD_INDEX_ARCH: dict[str, str] = {
"gfx908": "gfx908", # MI200/MI100
}
# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix
# (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every
# AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI.
# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix (bnb #1887); bnb <=0.49.2 NaNs on AMD. Drop once bnb 0.50+ ships on PyPI.
_BNB_ROCM_PRERELEASE_URLS: dict[str, str] = {
"x86_64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
@ -436,9 +419,7 @@ _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = {
"download/continuous-release_main/"
"bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl"
),
# Windows ROCm wheel ships libbitsandbytes_rocm{VER}.dll. BNB's HIP
# auto-detect may mismatch the DLL suffix, so we scan the wheel and set
# BNB_ROCM_VERSION in _install_bnb_windows_rocm() and worker.py.
# Windows ROCm wheel ships libbitsandbytes_rocm{VER}.dll; the wheel is scanned and BNB_ROCM_VERSION set to match.
"win_amd64": (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main/"
@ -475,8 +456,7 @@ def _path_inside_venv(path: str) -> bool:
try:
# realpath (not abspath): resolve symlinks/8.3 names so an aliased venv matches.
_root = os.path.normcase(os.path.realpath(sys.prefix))
# Guard a root-dir prefix (C:\ or /): commonpath would match every path on
# it. A venv is never at root, so treat that as outside.
# Root-dir prefix (C:\ or /) would commonpath-match everything; a venv is never at root.
if os.path.dirname(_root) == _root:
return False
return os.path.normcase(os.path.commonpath([os.path.realpath(path), _root])) == _root
@ -514,9 +494,7 @@ def _amd_smi_allowed() -> bool:
return True
if flag in ("0", "false", "no", "off"):
return False
# A real HIP SDK lets amd-smi run un-elevated; hipinfo-on-PATH is the proxy.
# Ignore the venv hipInfo.exe (AMD wheel via bnb fix): not a HIP SDK, doesn't
# stop amd-smi's DiskPart UAC.
# hipinfo-on-PATH proxies a real HIP SDK; the venv hipInfo.exe is not one.
if _external_hipinfo_on_path():
return True
for _var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"):
@ -539,16 +517,13 @@ def _detect_rocm_version() -> tuple[int, int] | None:
try:
with open(path, encoding = "utf-8") as fh:
parts = fh.read().strip().split("-")[0].split(".")
# Explicit length guard: don't rely on the broad except below to
# swallow IndexError on a single-component version (e.g. "6\n").
# Length guard for single-component versions (e.g. "6\n").
if len(parts) >= 2:
return int(parts[0]), int(parts[1])
except Exception:
pass
# Try amd-smi version (outputs "... | ROCm version: X.Y.Z").
# Gated off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt);
# hipconfig below covers that case.
# Try amd-smi version ("ROCm version: X.Y.Z"); gated off on Windows w/o a HIP SDK (UAC/DiskPart prompt).
amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None
if amd_smi:
try:
@ -585,10 +560,7 @@ def _detect_rocm_version() -> tuple[int, int] | None:
except Exception:
pass
# Distro package-manager fallbacks: package-managed ROCm can expose GPUs via
# rocminfo/amd-smi but lack /opt/rocm/.info/version and hipconfig, so probe
# dpkg (Debian/Ubuntu) and rpm (RHEL/Fedora/SUSE) for the rocm-core version.
# Matches install.sh::get_torch_index_url so `studio update` == fresh install.
# dpkg/rpm rocm-core fallback: package-managed ROCm may lack .info/version and hipconfig. Matches install.sh::get_torch_index_url.
for cmd in (
["dpkg-query", "-W", "-f=${Version}\n", "rocm-core"],
["rpm", "-q", "--qf", "%{VERSION}\n", "rocm-core"],
@ -660,8 +632,7 @@ def _detect_windows_gfx_arch() -> str | None:
def _dedup_pick(tokens: list[str]) -> "str | None":
if not tokens:
return None
# Index into the full ordered list so HIP_VISIBLE_DEVICES addresses
# GPU N on mixed-arch hosts, then return that arch.
# Index the full ordered list so HIP_VISIBLE_DEVICES addresses GPU N on mixed-arch hosts.
return tokens[_pick_visible_index(len(tokens))]
# 2. hipinfo via PATH, then HIP_PATH\bin / ROCM_PATH\bin.
@ -675,10 +646,7 @@ def _detect_windows_gfx_arch() -> str | None:
hipinfo = _candidate
break
if not hipinfo:
# 2b. AMD torch wheels ship hipInfo.exe into the venv Scripts dir
# (next to python.exe); resolvable even on driver-only hosts with no
# SDK install at all. Lets `studio update` re-detect the arch on a
# venv that already has the AMD wheel.
# 2b. AMD torch wheels ship hipInfo.exe into venv Scripts; lets `studio update` re-detect on driver-only hosts.
_venv_hipinfo = os.path.join(os.path.dirname(sys.executable), "hipInfo.exe")
if os.path.isfile(_venv_hipinfo):
hipinfo = _venv_hipinfo
@ -690,13 +658,9 @@ def _detect_windows_gfx_arch() -> str | None:
stderr = subprocess.DEVNULL,
timeout = 10,
)
# Accept partial output even when hipinfo crashes (e.g. 0xC0000005 /
# STATUS_ACCESS_VIOLATION on some RDNA 4 hosts): a gcnArchName in stdout
# means the device was enumerated pre-crash, so the arch is trustworthy.
# Ignoring it causes a silent CPU PyTorch fallback (issue #6043).
# Accept partial output even when hipinfo crashes (0xC0000005 on some RDNA 4, #6043): a pre-crash gcnArchName is trustworthy.
text = result.stdout.decode(errors = "replace")
# findall gets every gcnArchName line so multi-GPU hosts are
# enumerable and HIP_VISIBLE_DEVICES selects correctly.
# findall gets every gcnArchName line so HIP_VISIBLE_DEVICES selects on multi-GPU hosts.
_tokens = [
t.strip().lower() for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
]
@ -706,9 +670,7 @@ def _detect_windows_gfx_arch() -> str | None:
except Exception:
pass
# 3. amd-smi fallback -- runtime-only Radeon installs ship amd-smi but no hipinfo.
# Gated off on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt); the arch
# arrives via --rocm-gfx / name inference there, so this is only needed when safe.
# 3. amd-smi fallback (runtime-only Radeon installs lack hipinfo); gated off on Windows w/o a HIP SDK (UAC/DiskPart prompt).
amd_smi = shutil.which("amd-smi") if _amd_smi_allowed() else None
if amd_smi:
for _args in (("static", "--asic"), ("list",)):
@ -737,11 +699,7 @@ def _detect_windows_gfx_arch() -> str | None:
except Exception:
continue
# 4. Last resort: GPU marketing name via WMI → arch table. Driver-only
# hosts (Adrenalin, no HIP SDK) have neither hipinfo nor amd-smi
# (amd-smi does not exist on Windows at all), but the display driver
# always knows the GPU name. Mirrors setup.ps1's $nameArchTable so a
# standalone `studio update` can repair a CPU-only venv on such hosts.
# 4. Last resort: GPU marketing name via WMI → arch table (driver-only hosts have neither hipinfo nor amd-smi); mirrors setup.ps1's $nameArchTable.
try:
result = subprocess.run(
[
@ -771,10 +729,7 @@ def _detect_windows_gfx_arch() -> str | None:
return None
# GPU marketing-name → gfx arch table, mirroring setup.ps1's $nameArchTable.
# Most-specific first; first match wins. Covers only arches the ROCm
# prebuilts / AMD Windows torch indexes support; unknown names return None
# (callers then fall back cleanly to CPU).
# GPU marketing-name → gfx arch table (mirrors setup.ps1's $nameArchTable); most-specific first; unknown names return None (CPU fallback).
_WIN_GPU_NAME_ARCH_TABLE: "list[tuple[str, str]]" = [
(r"9070|9080", "gfx1201"), # RDNA 4 (Navi 48: Radeon RX 9070 XT / 9070 GRE / 9070 / 9080)
(r"9060", "gfx1200"), # RDNA 4 (Navi 44: Radeon RX 9060 XT / 9060)
@ -968,8 +923,7 @@ def _detect_bnb_rocm_dll_ver() -> str | None:
m = re.search(r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(dll))
if m:
all_vers.append(m.group(1))
# Highest numeric suffix wins (e.g. "713" over "72"); glob order is not
# guaranteed, so sort rather than take the first match.
# Highest numeric suffix wins ("713" over "72"); glob order is not guaranteed.
return max(all_vers, key = lambda v: int(v)) if all_vers else None
@ -1038,8 +992,7 @@ def _persist_bnb_rocm_version(version: str) -> bool:
existing = (
sitecustomize_path.read_text(encoding = "utf-8") if sitecustomize_path.exists() else ""
)
# Strip all managed regions, including one whose END marker was lost to
# an interrupted write, then append exactly one fresh block.
# Strip all managed regions (even END-marker-less from an interrupted write), append one fresh block.
pattern = re.compile(
rf"{re.escape(_BNB_ROCM_SITECUSTOMIZE_BEGIN)}.*?"
rf"(?:{re.escape(_BNB_ROCM_SITECUSTOMIZE_END)}\n?|\Z)",
@ -1079,10 +1032,7 @@ def _has_rocm_gpu() -> bool:
if _has_usable_nvidia_gpu():
return False
for cmd, check_fn in (
# rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit).
# gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines like
# "gfx11-generic"/"gfx9-4-generic" with only 1-2 digits before the dash,
# which must not be treated as a real GPU.
# rocminfo: real gfx GPU id only (gfx000 = CPU agent; generic "gfx11-generic" ISA lines are not GPUs).
(
["rocminfo"],
lambda out: bool(re.search(r"gfx[1-9][0-9a-z]{2,3}", out.lower())),
@ -1096,8 +1046,7 @@ def _has_rocm_gpu() -> bool:
exe = shutil.which(cmd[0])
if not exe:
continue
# Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt);
# rely on rocminfo / the sysfs fallback there.
# Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart prompt).
if cmd[0] == "amd-smi" and not _amd_smi_allowed():
continue
try:
@ -1114,14 +1063,8 @@ def _has_rocm_gpu() -> bool:
if result.returncode == 0 and result.stdout.strip():
if check_fn(result.stdout):
return True
# sysfs KFD topology fallback (Linux only) -- matches install.sh's runtime-only
# detection. On minimal package-managed installs (no rocminfo / amd-smi), the
# kernel exposes AMD GPUs via /sys/class/kfd so `studio update` can still repair.
#
# Guard: reject any KFD node whose properties file reports a non-AMD vendor. The
# NVIDIA open kernel module (driver 560+) registers KFD nodes with a non-zero
# gpu_id and vendor_id 4318 (0x10DE), not the AMD 4098 (0x1002); without this
# check the fallback returns True on NVIDIA-only hosts, installing ROCm wheels.
# sysfs KFD topology fallback (Linux, matches install.sh): minimal installs lack rocminfo/amd-smi.
# Reject non-AMD vendors: the NVIDIA open kernel module also registers KFD nodes (vendor 0x10DE).
if sys.platform != "win32":
try:
kfd_nodes = "/sys/class/kfd/kfd/topology/nodes"
@ -1135,10 +1078,7 @@ def _has_rocm_gpu() -> bool:
continue
if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node
continue
# Require AMD vendor_id 4098 (0x1002). KFD properties files exist
# on every kernel exposing /sys/class/kfd, so a missing file means
# AMD ownership is unconfirmed -- skip the node rather than risk a
# false positive (e.g. NVIDIA open-driver KFD nodes lacking it).
# Require AMD vendor_id 4098 (0x1002); a missing properties file leaves AMD ownership unconfirmed -- skip.
props_path = os.path.join(kfd_nodes, entry, "properties")
try:
with open(props_path, encoding = "utf-8") as fh:
@ -1184,8 +1124,7 @@ def _has_usable_nvidia_gpu() -> bool:
return True
except Exception:
pass
# Fallback: the NVIDIA driver exposes one subdirectory per GPU under
# /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state.
# Fallback: /proc/driver/nvidia/gpus/ has one subdir per GPU regardless of nvidia-smi state.
if sys.platform != "win32":
try:
gpu_dir = "/proc/driver/nvidia/gpus"
@ -1265,10 +1204,7 @@ def _install_bnb_windows_rocm() -> bool:
)
if not _ok:
return False
# Detect the actual ROCm DLL suffix in the wheel and set BNB_ROCM_VERSION so bnb
# loads the right DLL regardless of torch.version.hip (the wheel may ship "72"
# while torch reports 7.13). The worker subprocess inherits it; fall back to "72"
# if detection fails (e.g. a no-op / dry-run install).
# Detect the ROCm DLL suffix and set BNB_ROCM_VERSION (wheel may ship "72" while torch reports 7.13); fall back to "72".
_env_ver = os.environ.get("BNB_ROCM_VERSION")
_env_is_persisted_default = (
os.environ.get(_BNB_ROCM_VERSION_SOURCE_ENV) == _BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE
@ -1283,11 +1219,7 @@ def _install_bnb_windows_rocm() -> bool:
_persist_detected_version = True
if _persist_detected_version:
_persist_bnb_rocm_version(_ver)
# Make hipInfo.exe (shipped into venv Scripts by the AMD torch wheel) resolvable
# via PATH for this process and every child python (import checks, precompile):
# bitsandbytes runs hipinfo.exe at import to detect the GPU arch and logs a scary
# (harmless) ERROR + WARNING when it is missing. Scripts is on PATH only for an
# activated venv, which neither Unsloth nor the installer's children ever do.
# Put venv Scripts (hipInfo.exe from the AMD torch wheel) on PATH: bnb probes hipinfo.exe at import and logs a scary (harmless) ERROR when missing.
_scripts_dir = os.path.dirname(sys.executable)
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")) and not shutil.which(
"hipinfo.exe"
@ -1366,8 +1298,7 @@ def _is_pip_rocm_family_leaf(leaf: str) -> bool:
rocm7.2-private) starts with "rocm" but is a custom pin the verbatim path owns, so
match EXACTLY. Mirrors install.sh / setup.ps1.
"""
# gfx must be followed by a digit (gfx90a, gfx1151, gfx120X-all): a gfx-prefixed
# custom leaf (gfx-private) is a verbatim pin, like rocm7.2-private.
# gfx must be followed by a digit; a gfx-private custom leaf is a verbatim pin.
return bool(re.fullmatch(r"rocm\d+(?:\.\d+)?", leaf)) or bool(re.match(r"gfx\d", leaf))
@ -1508,8 +1439,7 @@ def _ensure_cuda_torch() -> None:
Only repairs when torch actually links against HIP/ROCm. Healthy CUDA
torch and deliberate CPU-only torch are left untouched.
"""
# Respect install.sh's backend: only "" (standalone update) or "cuda" force CUDA
# wheels; "rocm"/"cpu"/unrecognised are deliberate.
# Respect install.sh's backend: only "" (standalone update) or "cuda" force CUDA wheels.
if _TORCH_BACKEND not in ("", "cuda"):
return
# An explicit unknown-family pin was applied VERBATIM at install time; leave it alone.
@ -1521,11 +1451,9 @@ def _ensure_cuda_torch() -> None:
# Never undo a deliberate ROCm install (setup.ps1 sets this marker).
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
return
# An explicit CUDA pin (headless / CI cross-install) commits to CUDA wheels and skips ALL
# GPU probing, so it clears both the CUDA_VISIBLE_DEVICES hide gate and the NVIDIA gate below.
# An explicit CUDA pin commits to CUDA wheels and skips ALL GPU probing gates below.
_cuda_pinned = _explicit_cuda_torch_index_url() is not None
# CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU; never force CUDA
# wheels over that unless a CUDA index is pinned.
# CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU; honour it unless a CUDA index is pinned.
_cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if not _cuda_pinned and _cvd is not None and _cvd.strip() in ("", "-1"):
return
@ -1533,9 +1461,7 @@ def _ensure_cuda_torch() -> None:
if not _cuda_pinned and not _has_usable_nvidia_gpu():
return
# Classify the installed torch: "hip" (ROCm poisoning signature), "cuda" (healthy),
# or "cpu". A non-zero exit means torch is missing/un-importable: without a pin the
# base install owns it, but a pinned CUDA index reinstalls it below.
# Classify installed torch: "hip" (ROCm poisoning signature), "cuda", or "cpu"; non-zero exit = missing/un-importable.
try:
probe = subprocess.run(
[
@ -1558,9 +1484,7 @@ def _ensure_cuda_torch() -> None:
except (OSError, subprocess.TimeoutExpired):
return
if probe.returncode != 0:
# torch present but can't import. Without a pin the base install owns it; but an
# explicit CUDA pin forces this pass (failed probe) and the base update won't
# reinstall an already-installed torch, so reinstall from the pin (self-resolving).
# Un-importable torch: only an explicit CUDA pin reinstalls here (the base update won't touch an installed torch).
if not _cuda_pinned:
return
index_url = _detect_cuda_torch_index_url()
@ -1588,9 +1512,7 @@ def _ensure_cuda_torch() -> None:
if not _marker_lines:
return
_marker, _, _installed_cu = _marker_lines[-1].partition("|")
# Reinstall CUDA torch on a ROCm build on an NVIDIA host (poisoning signature), or when a
# CUDA index is pinned but the venv has the wrong family (CPU or a different cuXXX). A
# healthy match, or a CPU wheel with no CUDA pin, is left alone.
# Reinstall on ROCm-on-NVIDIA poisoning or a pinned-CUDA family mismatch; healthy/CPU-no-pin left alone.
_pin = _explicit_torch_index_url()
_pin_leaf = _torch_index_leaf(_pin) if _pin else ""
_pinned_cuda = _is_cuda_family_leaf(_pin_leaf)
@ -1599,8 +1521,7 @@ def _ensure_cuda_torch() -> None:
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:
# Installed cuXXX differs from the pin. An untagged build (empty) counts too:
# the family can't be confirmed, so reinstall to enforce it (idempotent).
# Installed cuXXX differs from the pin; an untagged build counts too (family unconfirmed, idempotent).
_installed_desc = _installed_cu if _installed_cu else "an untagged CUDA build"
_why = f"torch is {_installed_desc} but the pinned CUDA index is {_pin_leaf}"
else:
@ -1639,8 +1560,7 @@ def _ensure_cpu_torch() -> None:
if pin is None:
return
# Classify the installed torch family. A non-zero exit means torch is missing or
# un-importable: the explicit CPU pin reinstalls it below.
# Classify the installed torch family; non-zero exit = missing/un-importable -> the CPU pin reinstalls below.
try:
probe = subprocess.run(
[
@ -1662,9 +1582,7 @@ def _ensure_cpu_torch() -> None:
except (OSError, subprocess.TimeoutExpired):
return
if probe.returncode != 0:
# torch present but can't import. The explicit CPU pin forces this pass (failed
# probe) and the base update won't reinstall an already-installed torch, so
# reinstall from the pin (self-resolving, no loop).
# Un-importable torch: reinstall from the explicit CPU pin (self-resolving, no loop).
_torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC
print(
f" torch cannot import but an explicit CPU index is pinned -- reinstalling "
@ -1694,8 +1612,7 @@ def _ensure_cpu_torch() -> None:
" torch is a GPU build but an explicit CPU index is pinned -- reinstalling "
f"CPU torch from {_strip_index_url_credentials(pin)}"
)
# Pin the supported torch<2.11 family (the /cpu index now serves 2.11+, so a bare
# trio could resolve out of range or ABI-mismatched).
# Pin the supported torch family (the /cpu index now serves 2.11+).
_torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC
pip_install(
"CPU torch repair",
@ -1720,15 +1637,13 @@ def _ensure_rocm_torch() -> None:
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
global _rocm_windows_torch_installed
# install.sh's resolved backend is authoritative: skip ROCm when it already chose a
# non-ROCm family (avoids re-detecting in a subprocess that may see a different env).
# install.sh's resolved backend is authoritative: skip ROCm for a non-ROCm family.
if _TORCH_BACKEND in ("cuda", "cpu"):
return
# An explicit unknown-family pin was applied VERBATIM at install time; leave it alone.
if _explicit_unknown_family_torch_index_url() is not None:
return
# setup.ps1 sets this after installing AMD wheels; skip only when torch is actually
# importable as ROCm (a wiped venv leaves a stale env-var that must not suppress it).
# setup.ps1's marker; trust it only when torch actually imports as ROCm (a wiped venv leaves a stale env-var).
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
_torch_ok = False
try:
@ -1752,8 +1667,7 @@ def _ensure_rocm_torch() -> None:
pass
if _torch_ok:
_rocm_windows_torch_installed = True
# ROCm torch is already installed, but the AMD Windows BNB wheel is still
# needed (the PyPI bitsandbytes ships only CUDA DLLs, fails on ROCm).
# AMD Windows BNB wheel still needed (PyPI bitsandbytes ships only CUDA DLLs).
_install_bnb_windows_rocm()
return
# torch was wiped between runs; fall through to the full install path
@ -1761,10 +1675,7 @@ def _ensure_rocm_torch() -> None:
return
if IS_WINDOWS:
# An explicit ROCm-family pin commits to ROCm wheels regardless of the visible
# GPU and overrides the public per-arch index (mirrors the Linux pin handling
# below): after a pinned setup.ps1 install fails to CPU, this repair must retry
# the PINNED index, not repo.amd.com.
# An explicit ROCm-family pin commits to ROCm wheels and overrides the public per-arch index: retry the PINNED index, not repo.amd.com.
_win_rocm_pin = _explicit_rocm_torch_index_url()
if _win_rocm_pin is None and _has_usable_nvidia_gpu():
return
@ -1802,14 +1713,11 @@ def _ensure_rocm_torch() -> None:
f" {gfx_arch or 'pinned ROCm index'} (Windows) -- installing torch from "
f"{_strip_index_url_credentials(index_url)}"
)
# Pin companions for the arches install.ps1/setup.ps1 pin (gfx120X / Strix)
# so the per-arch index resolves an ABI-consistent trio; other arches stay bare.
# Pin companions for the arches install.ps1/setup.ps1 pin so the per-arch index resolves an ABI-consistent trio.
_torch_pkg, _vision_pkg, _audio_pkg = _WINDOWS_ROCM_TORCH_PKG_SPECS.get(
gfx_arch, ("torch", "torchvision", "torchaudio")
)
# Nonfatal: a transient AMD-index failure must not abort the install.
# --force-reinstall resolves before uninstalling, so a failed index keeps the
# existing build intact; let the user retry.
# Nonfatal: --force-reinstall resolves before uninstalling, so a failed index keeps the existing build.
if not pip_install_try(
f"ROCm torch (Windows, {gfx_arch or 'pinned'})",
"--force-reinstall",
@ -1826,14 +1734,9 @@ def _ensure_rocm_torch() -> None:
"later to retry ROCm."
)
return
# ROCm torch is installed (or already was); flag it so later phases
# do not overwrite it with the generic CPU torch wheel. BNB is a
# separate dependency -- a BNB install failure must NOT roll back the
# torch ROCm install.
# Flag ROCm torch installed so later phases don't overwrite it; a BNB failure must NOT roll it back.
_rocm_windows_torch_installed = True
# Always install AMD Windows bitsandbytes -- the PyPI wheel ships only
# CUDA DLLs and fails on ROCm. Install even when torch was already a
# ROCm build so `studio update` repairs a broken bnb.
# Always install AMD Windows bitsandbytes (PyPI wheel ships only CUDA DLLs); also repairs a broken bnb on update.
if not _install_bnb_windows_rocm():
print(
" Warning: AMD Windows bitsandbytes install failed; "
@ -1844,8 +1747,7 @@ 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
# An explicit ROCm pin commits to ROCm wheels regardless of the visible GPU (headless / CI).
# Mirror _ensure_cuda_torch: skip the NVIDIA/no-AMD/unreadable gates.
# An explicit ROCm pin commits to ROCm wheels regardless of the visible GPU (headless / CI); skip the GPU gates.
_rocm_pin = _explicit_rocm_torch_index_url()
_inferred_linux_gfx = (
_infer_linux_amd_gfx_arch() if (_rocm_pin is None and not IS_WINDOWS) else None
@ -1867,9 +1769,7 @@ def _ensure_rocm_torch() -> None:
# Explicit pin or inferred gfx: the index drives the install.
ver = (0, 0)
# Probe whether torch links against HIP, capturing the installed ROCm tag for pin-mismatch
# detection. Emit ONE "<hip_marker>|<version>" line: marker (HIP version, "rocm" sentinel,
# or empty for CPU/CUDA) before "|", wheel version after.
# Probe HIP linkage; emit ONE "<hip_marker>|<version>" line for pin-mismatch detection.
try:
probe = subprocess.run(
[
@ -1879,8 +1779,7 @@ def _ensure_rocm_torch() -> None:
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"ver=getattr(torch,'__version__','').lower(); "
# HIP version if present, else a "rocm" sentinel when only the
# version string flags ROCm; empty marker = CPU/CUDA torch.
# HIP version, "rocm" sentinel, or empty marker = CPU/CUDA torch.
"marker=hip if hip else ('rocm' if 'rocm' in ver else ''); "
"print(marker + '|' + ver)"
),
@ -1903,9 +1802,7 @@ def _ensure_rocm_torch() -> None:
# A "|"-delimited line is required; without it treat HIP as absent -> reinstall.
has_hip_torch = bool(_sep) and _hip_marker != ""
# An explicit ROCm pin whose family differs from the installed torch must reinstall, else a
# rocm7.2/gfx* pin over an older +rocm6.4/7.1 build never applies. Version-tag heuristic
# only: a same-tag per-arch switch (gfx1151 -> gfx120X-all, both +rocm7.13.0) isn't detectable.
# A ROCm pin whose family differs from the installed torch must reinstall; version-tag heuristic only (same-tag per-arch switch undetectable).
_rocm_pin_mismatch = (
_rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
if (has_hip_torch and _rocm_pin is not None)
@ -1975,8 +1872,7 @@ def _ensure_rocm_torch() -> None:
_strix_gfx = {"gfx1151", "gfx1150", "gfx1152"}
_detected_strix = _strix_gfx.intersection(gfx_codes)
if _detected_strix:
# Runtime-visible GPU (HIP_VISIBLE_DEVICES index into gfx_codes, else first);
# skip the override unless it's Strix.
# Runtime-visible GPU (HIP_VISIBLE_DEVICES index, else first) must be Strix.
_runtime_gfx = gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None
if _runtime_gfx in _strix_gfx:
_selected_gfx = _runtime_gfx
@ -1986,8 +1882,7 @@ def _ensure_rocm_torch() -> None:
_strix_override_url = f"{_amd_mirror}/{_selected_gfx}/"
_strix_override_pkgs = (
"torch>=2.11.0,<2.12.0",
# Pin companions to the 2.11.x range: the exclusive --index-url could
# otherwise resolve a build for a different torch major (ABI mismatch).
# Pin companions to 2.11.x (exclusive --index-url could resolve ABI-mismatched).
"torchvision>=0.26.0,<0.27.0",
"torchaudio>=2.11.0,<2.12.0",
)
@ -2035,8 +1930,7 @@ def _ensure_rocm_torch() -> None:
f" requires a source build of bitsandbytes for gfx906 (see docs.unsloth.ai/amd).\n"
)
# The Strix override must fire even when has_hip_torch is True: an existing
# torch.version.hip == "7.1" is exactly the broken combo it repairs.
# Strix override fires even when has_hip_torch: hip == "7.1" is exactly the broken combo.
if _strix_override_url is not None and _strix_override_pkgs is not None:
index_url = _strix_override_url
_torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs
@ -2104,8 +1998,7 @@ def _ensure_rocm_torch() -> None:
if _override_idx is None:
index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
print(f" ROCm torch -- installing from {_strip_index_url_credentials(index_url)}")
# Only the _grouped_mm-bug gfx arches need the 2.11 spec; other gfx indexes ship
# <2.11 and stay on the default range (matches install.ps1 / setup.ps1).
# Only the _grouped_mm-bug gfx arches need the 2.11 spec (matches install.ps1 / setup.ps1).
if tag in _ROCM_GFX_TORCH211_LEAVES:
_torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["rocm7.2"]
elif tag.startswith("gfx"):
@ -2149,10 +2042,7 @@ def _ensure_rocm_torch() -> None:
[sys.executable, "-m", "pip", "uninstall", "-y", "bitsandbytes"],
capture_output = True,
)
# Install bitsandbytes only when torch links against ROCm. Prefers the
# continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix), falling back
# to PyPI when the pre-release wheel won't install. Use pip for the
# pre-release wheel because uv rejects its filename/metadata version mismatch.
# bitsandbytes only when torch links ROCm; prefer the pre-release wheel (bnb #1887), pip not uv (filename/metadata version mismatch).
elif rocm_torch_ready:
_bnb_url = _bnb_rocm_prerelease_url()
_bnb_installed = False
@ -2184,9 +2074,6 @@ def _ensure_rocm_torch() -> None:
)
# _uv_safe_path is imported from backend.utils.uv_path_safety (shared with mlx_repair).
def _windows_hidden_subprocess_kwargs() -> dict[str, object]:
"""Return Windows-only subprocess kwargs that suppress console windows."""
if not IS_WINDOWS:
@ -2224,11 +2111,9 @@ def _infer_no_torch() -> bool:
NO_TORCH = _infer_no_torch()
# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() ("cuda", "rocm",
# "cpu"; empty = standalone `studio update`, where we re-detect).
# UNSLOTH_TORCH_BACKEND is set by install.sh ("cuda"/"rocm"/"cpu"; empty = standalone `studio update`).
_TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower()
# Standalone update with an explicit pin: derive the backend from the override (classify on
# the final URL/family segment, mirroring install.sh) instead of re-probing the GPU.
# Standalone update with an explicit pin: derive the backend from the override leaf (mirrors install.sh).
if not _TORCH_BACKEND:
_idx_override = (
os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
@ -2240,9 +2125,7 @@ if not _TORCH_BACKEND:
elif _idx_leaf == "cpu":
_TORCH_BACKEND = "cpu"
elif _is_cuda_family_leaf(_idx_leaf):
# Require a digit after "cu" so /current or /custom is NOT branded CUDA (a wrong backend
# makes _ensure_rocm_torch return early on AMD hosts). An unknown leaf keeps "" so the
# helpers probe the GPU.
# Require a digit after "cu" so /current or /custom is NOT branded CUDA; an unknown leaf keeps "" (helpers probe the GPU).
_TORCH_BACKEND = "cuda"
@ -2264,15 +2147,10 @@ def _torch_step_label(suffix: str) -> str:
# -- Verbosity control ----------------------------------------------------------
# By default the installer shows a minimal in-place one-line progress bar.
# Set UNSLOTH_VERBOSE=1 to restore full per-step output:
# CLI: unsloth studio setup --verbose
# Linux/Mac: UNSLOTH_VERBOSE=1 ./studio/setup.sh
# Windows: $env:UNSLOTH_VERBOSE="1" ; .\studio\setup.ps1
# Default: minimal in-place progress bar; UNSLOTH_VERBOSE=1 restores full per-step output.
VERBOSE: bool = os.environ.get("UNSLOTH_VERBOSE", "0") == "1"
# Progress bar state -- updated by _progress() per install step.
# Update _TOTAL if you add/remove steps in install_python_stack().
# Progress bar state -- update _TOTAL if you add/remove steps in install_python_stack().
_STEP: int = 0
_TOTAL: int = 0 # set at runtime in install_python_stack() based on platform
_PROGRESS_LINE_ACTIVE: bool = False
@ -2287,20 +2165,16 @@ LOCAL_DD_UNSTRUCTURED_PLUGIN = (
)
LOCAL_DD_GITHUB_PLUGIN = SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed"
# mlx-lm 0.31.3 broke gemma4 / qwen3_5 loading (strict load_weights rejects the
# QK-norm q_norm/k_norm tensors); exclude just that release. See mlx-lm #1242.
# mlx-lm 0.31.3 broke gemma4 / qwen3_5 QK-norm loading; exclude just that release (mlx-lm #1242).
MLX_LM_BAD_VERSION_EXCLUSION = "!=0.31.3"
# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides).
# _uv_safe_path: uv truncates UV_OVERRIDE at the first space too (issue #6503).
# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin; _uv_safe_path because uv truncates UV_OVERRIDE at the first space (#6503).
_MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt"
if IS_MAC_ARM and _MLX_OVERRIDES.is_file() and "UV_OVERRIDE" not in os.environ:
os.environ["UV_OVERRIDE"] = _uv_safe_path(_MLX_OVERRIDES)
# -- Unicode-safe printing ---------------------------------------------
# On Windows the console encoding may be a legacy code page (e.g. CP1252)
# that cannot represent glyphs like ✅ or ❌. _safe_print() degrades to ASCII
# equivalents so the installer never crashes over a status glyph.
# Windows console may be a legacy code page (e.g. CP1252); _safe_print() degrades glyphs to ASCII.
_UNICODE_TO_ASCII: dict[str, str] = {
"\u2705": "[OK]", # ✅
@ -2362,8 +2236,7 @@ def _stdout_supports_color() -> bool:
_HAS_COLOR = _stdout_supports_color()
# Column layout — matches setup.sh step() helper:
# 2-space indent, 15-char label (dim), then value.
# Column layout — matches setup.sh step(): 2-space indent, 15-char dim label, then value.
_LABEL = "deps"
_COL = 15
_INDENT = 2
@ -2465,8 +2338,7 @@ def run(
if result.returncode != 0:
_step("error", f"{label} failed (exit code {result.returncode})", _red)
if result.stdout:
# Redact before printing: the failing pip command may carry a pinned --index-url
# with userinfo/?token= creds, so raw pip error text would leak them.
# Redact before printing: pip error text may embed a pinned --index-url's userinfo/?token= creds.
print(_redact_install_output(result.stdout))
sys.exit(result.returncode)
return result
@ -2475,13 +2347,8 @@ def run(
# Packages to skip on Windows (require special build steps)
WINDOWS_SKIP_PACKAGES = {"triton_kernels"}
# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode). These
# either *are* torch extensions or have unconditional ``Requires-Dist: torch``, so
# installing them pulls torch back in. ``librosa`` is here despite not requiring
# torch: upstream ``llvmlite`` dropped its macOS x86_64 wheel (0.46.0+ ships only
# macosx_arm64 / manylinux / win_amd64), so on Intel Mac the librosa -> numba ->
# llvmlite chain triggers a from-source build that fails without LLVM 14/15 headers.
# Tracked in unslothai/unsloth#5046.
# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode): torch extensions or hard
# ``Requires-Dist: torch``; ``librosa`` because its numba -> llvmlite chain fails from-source on Intel Mac (#5046).
NO_TORCH_SKIP_PACKAGES = {
"torch-stoi",
"timm",
@ -2570,8 +2437,7 @@ def _bootstrap_uv() -> bool:
global UV_NEEDS_SYSTEM
if not shutil.which("uv"):
return False
# Probe: try a dry-run install targeting the current Python explicitly.
# Without --python, uv can ignore the activated venv on some platforms.
# Dry-run probe with explicit --python: uv can ignore the activated venv on some platforms.
probe = subprocess.run(
["uv", "pip", "install", "--dry-run", "--python", sys.executable, "pip"],
stdout = subprocess.PIPE,
@ -2645,25 +2511,18 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
cmd = ["uv", "pip", "install"]
if UV_NEEDS_SYSTEM:
cmd.append("--system")
# Always pass --python so uv targets the right environment. Without it, uv
# can ignore an activated venv and install into the system Python (seen on
# Colab and similar).
# Always pass --python so uv targets the right env (uv can ignore an activated venv, e.g. Colab).
cmd.extend(["--python", sys.executable])
cmd.extend(_translate_pip_args_for_uv(args))
# Torch is pre-installed, so don't add --torch-backend by default (solver dead-ends on
# CPU-only machines); callers can set UV_TORCH_BACKEND. Never add it to a pinned-index
# command: uv's torch backend redirects torch to its own per-backend index, defeating the pin.
# No --torch-backend by default (torch pre-installed); never on a pinned-index command (UV_TORCH_BACKEND redirects torch, defeating the pin).
_tb = os.environ.get("UV_TORCH_BACKEND", "")
if _tb and not _is_pinned_index_cmd(cmd):
cmd.append(f"--torch-backend={_tb}")
return cmd
# uv resolves --index-url / --default-index at LOWEST priority, so an inherited UV_INDEX /
# UV_EXTRA_INDEX_URL mirror wins and a pinned torch repair silently ignores the pin.
# Neutralise these for pinned installs (as install.sh #6898 / install.ps1 / setup.ps1 do).
# UV_TORCH_BACKEND redirects torch; PIP_* matter for the pip FALLBACK; UV_CONFIG_FILE is
# stripped + UV_NO_CONFIG=1 (a discovered uv.toml outranks the CLI pin, uv 0.10).
# uv resolves --index-url at LOWEST priority, so inherited index env vars silently defeat a pinned
# torch repair; neutralise these for pinned installs. UV_CONFIG_FILE is stripped + UV_NO_CONFIG=1.
_UV_INDEX_ENV_VARS = (
"UV_CONFIG_FILE",
"UV_DEFAULT_INDEX",
@ -2674,8 +2533,7 @@ _UV_INDEX_ENV_VARS = (
"UV_FIND_LINKS",
"PIP_EXTRA_INDEX_URL",
"PIP_FIND_LINKS",
# PIP_NO_INDEX=1 makes the pip fallback ignore ALL indexes (defeating --index-url);
# PIP_INDEX_URL is dropped too so a stale mirror env can't outrank the pin.
# PIP_NO_INDEX would defeat --index-url; PIP_INDEX_URL dropped so a stale mirror can't outrank the pin.
"PIP_NO_INDEX",
"PIP_INDEX_URL",
)
@ -2839,9 +2697,7 @@ def install_python_stack() -> int:
global USE_UV, _STEP, _TOTAL
_STEP = 0
# install.sh sets SKIP_STUDIO_BASE=1 to avoid reinstalling base packages;
# `studio update` does NOT, so unsloth + unsloth-zoo are reinstalled to pick
# up new versions.
# install.sh sets SKIP_STUDIO_BASE=1; `studio update` does NOT, so base packages are reinstalled.
skip_base = os.environ.get("SKIP_STUDIO_BASE", "0") == "1"
# --package installs a different package name (for testing).
package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth")
@ -2856,8 +2712,7 @@ def install_python_stack() -> int:
base_total += 2 # flash-attn + torch final repair (step 13), Linux
_TOTAL = (base_total - 1) if skip_base else base_total
# 1. Try uv for faster installs (before pip upgrade -- uv venvs don't
# include pip by default).
# 1. Try uv for faster installs (before pip upgrade -- uv venvs omit pip).
USE_UV = _bootstrap_uv()
# 2. Ensure pip is available (uv venvs from install.sh omit pip).
@ -2875,8 +2730,7 @@ def install_python_stack() -> int:
],
)
else:
# pip may not exist yet (uv-created venvs omit it). Try ensurepip,
# then upgrade. Direct upgrade only when pip is already present.
# pip may not exist yet (uv-created venvs omit it): ensurepip, else direct upgrade.
_has_pip = (
subprocess.run(
[sys.executable, "-m", "pip", "--version"],
@ -2898,10 +2752,7 @@ def install_python_stack() -> int:
[sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
)
# macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the
# mlx-vlm / mlx-lm transformers pin -- set at module load).
# Exclude mlx-lm 0.31.3 (see MLX_LM_BAD_VERSION_EXCLUSION); it broke
# gemma4 / qwen3_5 QK-norm loading. mlx-lm #1242.
# macOS arm64: MLX stack at latest (UV_OVERRIDE relaxes the transformers pin); exclude mlx-lm 0.31.3 (mlx-lm #1242).
if IS_MAC_ARM and not skip_base:
_progress("MLX stack (Apple Silicon)")
pip_install(
@ -2925,8 +2776,7 @@ def install_python_stack() -> int:
if skip_base:
pass
elif NO_TORCH:
# No-torch update path: install unsloth + unsloth-zoo, then runtime deps,
# both with --no-deps (PyPI metadata declares torch a hard dep; avoid it).
# No-torch update path: --no-deps throughout (PyPI metadata declares torch a hard dep).
_progress("base packages (no torch)")
pip_install(
f"Updating {package_name} + unsloth-zoo (no-torch mode)",
@ -2939,9 +2789,7 @@ def install_python_stack() -> int:
package_name,
"unsloth-zoo",
)
# Resolve pydantic WITH deps so pip pins pydantic-core to the exact version
# its metadata declares (under --no-deps pip picks the latest of each and
# trips pydantic's _ensure_pydantic_core_version check). Deps are torch-free.
# pydantic WITH deps so pip pins a matching pydantic-core (--no-deps trips _ensure_pydantic_core_version). Deps are torch-free.
pip_install(
"Installing pydantic (with deps for compatible core)",
"--no-cache-dir",
@ -2973,8 +2821,7 @@ def install_python_stack() -> int:
constrain = False,
)
elif local_repo:
# Local dev install: update deps from base.txt, then overlay the local
# checkout as an editable install (--no-deps so torch is not re-resolved).
# Local dev install: update deps, then overlay the local checkout editable (--no-deps).
_progress("base packages")
pip_install(
"Updating base packages",
@ -3012,9 +2859,7 @@ def install_python_stack() -> int:
package_name,
)
else:
# Update path: upgrade only unsloth + unsloth-zoo, preserving existing
# torch/CUDA installs. Torch is pre-installed by install.sh/setup.ps1;
# --upgrade-package targets only base pkgs.
# Update path: upgrade only unsloth + unsloth-zoo, preserving the pre-installed torch.
_progress("base packages")
pip_install(
"Updating base packages",
@ -3026,17 +2871,14 @@ def install_python_stack() -> int:
req = REQ_ROOT / "base.txt",
)
# 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the
# venv got CPU-only torch (common when pip resolves torch from PyPI).
# Must follow base packages so torch is present for inspection.
# 2b. Torch repair (wrong-family / CPU-only torch); must follow base packages so torch is present.
if not IS_MACOS and not NO_TORCH:
_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).
# Windows + AMD GPU: warn if ROCm torch was not installed.
if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu():
# Validate actual AMD GPU presence (not just tool existence).
import re as _re_win
@ -3052,10 +2894,7 @@ def install_python_stack() -> int:
_wexe = shutil.which(_wcmd[0])
if not _wexe:
continue
# Skip amd-smi on Windows w/o a HIP SDK (avoids the UAC/DiskPart
# prompt), as _has_rocm_gpu()/_detect_amd_gfx_codes do. The only loss
# is the best-effort "AMD GPU detected" note; ROCm-torch state below
# comes from the install itself.
# Skip amd-smi w/o a HIP SDK (UAC/DiskPart prompt); only loss is the best-effort note.
if _wcmd[0] == "amd-smi" and not _amd_smi_allowed():
continue
try:
@ -3099,16 +2938,11 @@ def install_python_stack() -> int:
req = REQ_ROOT / "extras-no-deps.txt",
)
# 4. Overrides (torchao) -- force-reinstall to a version matching the venv's
# torch so its C++ extensions load (see _select_torchao_spec). Skipped when
# torch is unavailable (Intel Mac GGUF-only) and on Windows ROCm (no working
# build; see below).
# 4. Overrides (torchao) -- force-reinstall to match the venv's torch (see _select_torchao_spec); skipped for no-torch and Windows ROCm.
if NO_TORCH:
_progress("dependency overrides (skipped, no torch)")
elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm():
# No working Windows ROCm torchao build: it imports an absent c10d backend
# and crashes transformers.quantizers. Unsloth stubs it at runtime, so
# installing it only ships a package that crashes on import -- skip it.
# No working Windows ROCm torchao build (crashes on import; stubbed at runtime) -- skip it.
_progress("dependency overrides (skipped, Windows ROCm)")
_safe_print(" Windows ROCm -- skipping torchao (no working build; stubbed at runtime)")
else:
@ -3123,8 +2957,7 @@ def install_python_stack() -> int:
_torchao_spec,
)
# 5. Triton kernels (no-deps, from source). Skip on Windows and macOS
# (no support).
# 5. Triton kernels (no-deps, from source); skip on Windows and macOS.
if not IS_WINDOWS and not IS_MACOS:
_progress("triton kernels")
pip_install(

File diff suppressed because it is too large Load diff

View file

@ -437,27 +437,30 @@ class TestKnown211SetParity:
), f"{label} floor gate must not use the unanchored ^rocm(\\d+)\\.(\\d+) prefix"
def test_install_ps1_bounds_unknown_leaf_pinned_torch(self):
"""install.ps1's pinned-torch install must bound BOTH companions on EVERY
index, cu<digits> families included: torchaudio 2.11 dropped its exact torch
pin from the wheel metadata, so a bare companion beside torch<2.11 can
resolve a mismatched 2.11.0 build (Codex P2, then unconditional per the
torchaudio 2.11 unpinning)."""
"""install.ps1's pinned-torch install must bound the whole trio on EVERY
index with the default torch 2.11 line (<2.12 trio, matching install.sh's
ceiling-composed default and _CUDA_TORCH_PKG_SPEC): torchaudio 2.11
dropped its exact torch pin from the wheel metadata, so a bare companion
beside a capped torch can resolve a mismatched build."""
text = INSTALL_PS1.read_text(encoding = "utf-8")
assert (
'$_pinVisionSpec = "torchvision>=0.19,<0.26.0"' in text
), "install.ps1 custom-pin install must bound torchvision (>=0.19,<0.26.0)"
'$_pinTorchSpec = "torch>=2.4,<2.12.0"' in text
), "install.ps1 default install must use the torch 2.11 line (<2.12.0)"
assert (
'$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"' in text
), "install.ps1 custom-pin install must bound torchaudio (>=2.4,<2.11.0)"
# No cu-family exemption: the bounds apply unconditionally.
'$_pinVisionSpec = "torchvision>=0.19,<0.27.0"' in text
), "install.ps1 must pair torchvision <0.27.0 with torch <2.12"
assert (
"$_pinCuLeaf" not in text
), "install.ps1 must bound companions on every index (no cu-family exemption)"
# The bounded companions must actually be passed to the install command.
'$_pinAudioSpec = "torchaudio>=2.4,<2.12.0"' in text
), "install.ps1 must pair torchaudio <2.12.0 with torch <2.12"
# No stale 2.10-line default remains anywhere in the Windows installer.
assert (
'"torch>=2.4,<2.11.0"' not in text
), "install.ps1 must not retain a <2.11.0 default torch range"
# The bounded trio must actually be passed to the install command.
assert re.search(
r'"torch>=2\.4,<2\.11\.0" \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl',
r"\$_pinTorchSpec \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl",
text,
), "install.ps1 custom-pin install must pass the bounded companion specs to uv"
), "install.ps1 pinned install must pass the bounded trio specs to uv"
def test_gfx_allowlist_matches_across_installers(self):
# The gfx 2.11 allowlist {gfx120x-all, gfx1151, gfx1150} must appear in each.
@ -696,9 +699,9 @@ class TestPinnedIndexClearsUvEnvParity:
# The custom-leaf branch bounds torch AND both companions (parity with the
# other installers' custom-pin trio bounds), gated on a non-cu-family leaf.
for spec in (
'$cudaTorchSpec = "torch>=2.4,<2.11.0"',
'$cudaVisionSpec = "torchvision>=0.19,<0.26.0"',
'$cudaAudioSpec = "torchaudio>=2.4,<2.11.0"',
'$cudaTorchSpec = "torch>=2.4,<2.12.0"',
'$cudaVisionSpec = "torchvision>=0.19,<0.27.0"',
'$cudaAudioSpec = "torchaudio>=2.4,<2.12.0"',
):
assert spec in text, f"setup.ps1 must bound the custom-leaf trio: {spec}"
assert (

View file

@ -64,45 +64,46 @@ class TestStructuralTorchConstraint:
_sh = _read(_INSTALL_SH)
def test_default_assignment_exists(self):
assert 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' in self._sh
"""The default range composes the per-file ceiling variable, so the
supported line (torch 2.11 today) is bumped in one place."""
assert '_TORCH_CEILING="2.12.0"' in self._sh
assert 'TORCH_CONSTRAINT="torch>=2.4,<${_TORCH_CEILING}"' in self._sh
def test_tightened_assignment_exists(self):
assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh
assert 'TORCH_CONSTRAINT="torch>=2.6,<${_TORCH_CEILING}"' 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_companion_ceilings_composed(self):
"""Companions bound to the same window via their own ceiling vars."""
assert '_TORCHVISION_CEILING="0.27.0"' in self._sh
assert '_TORCHAUDIO_CEILING="2.12.0"' in self._sh
assert 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<${_TORCHVISION_CEILING}"' in self._sh
assert 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<${_TORCHAUDIO_CEILING}"' in self._sh
def test_variable_used_in_pip_install(self):
"""$TORCH_CONSTRAINT must appear in a uv pip install line."""
assert '"$TORCH_CONSTRAINT"' in self._sh
def test_hardcoded_torch_constraint_only_on_assignments(self):
"""The hard-coded torch>=2.4,<2.11.0 string must only appear on
TORCH_CONSTRAINT= assignment lines, never on a pip/uv install line
(those must reference $TORCH_CONSTRAINT). Two assignments are expected:
the default, and the gfx906 (MI50) reroute that restores the default
<2.11 window after the rocm7.2 floor bump raised it to 2.11."""
hits = [ln for ln in self._sh.splitlines() if '"torch>=2.4,<2.11.0"' in ln]
assert hits, "default constraint literal missing from install.sh"
for ln in hits:
assert (
"TORCH_CONSTRAINT=" in ln
), f"torch>=2.4,<2.11.0 hardcoded off a TORCH_CONSTRAINT= assignment: {ln.strip()!r}"
assert (
"pip install" not in ln
), f"torch>=2.4,<2.11.0 hardcoded on a pip install line: {ln.strip()!r}"
"""The default range is composed from the ceiling vars, so the supported
line is bumped in one place. A hard-coded range may still appear on a
curated per-index TORCH_CONSTRAINT= override -- the gfx906 (MI50) reroute
caps below 2.11 because the rocm6.3 index tops out at torch 2.9.x -- but
never on a pip/uv install line (those must reference $TORCH_CONSTRAINT)."""
for literal in ('"torch>=2.4,<2.11.0"', '"torch>=2.4,<2.12.0"'):
for ln in self._sh.splitlines():
if literal not in ln:
continue
assert (
"TORCH_CONSTRAINT=" in ln
), f"{literal} hardcoded off a TORCH_CONSTRAINT= assignment: {ln.strip()!r}"
assert (
"pip install" not in ln
), f"{literal} hardcoded on a pip install line: {ln.strip()!r}"
def test_gfx906_reroute_caps_below_211(self):
"""The gfx906 / MI50 reroute must keep its literal sub-2.11 cap: the
rocm6.3 index it routes to serves no torch 2.11 wheel."""
assert self._sh.count('TORCH_CONSTRAINT="torch>=2.4,<2.11.0"') == 1
def test_tightening_guarded_by_skip_torch(self):
"""The block must check SKIP_TORCH=false."""
@ -132,7 +133,7 @@ class TestStructuralInstallPs1Unchanged:
assert "$TorchConstraint" not in self._ps1
def test_hardcoded_torch_constraint_present(self):
assert '"torch>=2.4,<2.11.0"' in self._ps1
assert '"torch>=2.4,<2.12.0"' in self._ps1
class TestInstallPs1UvDefaultIndex:

View file

@ -571,7 +571,7 @@ echo "=== Apple Silicon x86_64 (Rosetta) venv rebuild ==="
# Extract the real guard block from install.sh so we exercise the shipped logic
# (comment header down to its column-0 closing fi).
_GUARD_FILE=$(mktemp)
awk '/Guard against two independent Apple Silicon venv problems/{f=1} f{print} f&&/^fi$/{exit}' \
awk '/independent Apple Silicon venv/{f=1} f{print} f&&/^fi$/{exit}' \
"$INSTALL_SH" > "$_GUARD_FILE"
if [ ! -s "$_GUARD_FILE" ]; then

View file

@ -76,11 +76,12 @@ run_constraint_snippet() {
OS=\"$_os\"
_ARCH=\"$_arch\"
VENV_DIR=\"$_venv_dir\"
TORCH_CONSTRAINT=\"torch>=2.4,<2.11.0\"
_TORCH_CEILING=\"2.12.0\"
TORCH_CONSTRAINT=\"torch>=2.4,<\${_TORCH_CEILING}\"
if [ \"\$SKIP_TORCH\" = false ] && [ \"\$OS\" = \"macos\" ] && [ \"\$_ARCH\" = \"arm64\" ]; then
_PY_MINOR=\$(\"\$VENV_DIR/bin/python\" -c \"import sys; print(sys.version_info.minor)\" 2>/dev/null || echo \"0\")
if [ \"\$_PY_MINOR\" -ge 13 ] 2>/dev/null; then
TORCH_CONSTRAINT=\"torch>=2.6,<2.11.0\"
TORCH_CONSTRAINT=\"torch>=2.6,<\${_TORCH_CEILING}\"
fi
fi
echo \"\$TORCH_CONSTRAINT\"
@ -94,13 +95,25 @@ echo "=== Structural: TORCH_CONSTRAINT in install.sh ==="
_SH_CONTENT=$(cat "$INSTALL_SH")
# The supported line is centralized in per-file ceiling variables so a future
# torch 2.12 bump is a three-line change; the default range admits torch 2.11.
_count=$(grep -c '_TORCH_CEILING="2.12.0"' "$INSTALL_SH" || true)
assert_eq "torch ceiling variable defined once" "1" "$_count"
_count=$(grep -c '_TORCHVISION_CEILING="0.27.0"' "$INSTALL_SH" || true)
assert_eq "torchvision ceiling variable defined once" "1" "$_count"
_count=$(grep -c '_TORCHAUDIO_CEILING="2.12.0"' "$INSTALL_SH" || true)
assert_eq "torchaudio ceiling variable defined once" "1" "$_count"
# Each hardware branch assigns its own triple, so counting every occurrence made
# adding a branch (gfx906 in #7354) a test edit. The default is the one assigned at
# top level; a branch's is always indented, so anchor on that instead of counting.
_count=$(grep -c '^TORCH_CONSTRAINT="torch>=2.4,<2.11.0"$' "$INSTALL_SH" || true)
assert_eq "default TORCH_CONSTRAINT assignment exists" "1" "$_count"
# The default no longer spells the ceiling out, it composes _TORCH_CEILING, so the
# column-0 anchor goes on the composed form -- an indented branch pin, literal or
# composed, still cannot satisfy it.
_count=$(grep -c '^TORCH_CONSTRAINT="torch>=2.4,<${_TORCH_CEILING}"$' "$INSTALL_SH" || true)
assert_eq "default TORCH_CONSTRAINT assignment exists at top level" "1" "$_count"
_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' "$INSTALL_SH" || true)
_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.6,<${_TORCH_CEILING}"' "$INSTALL_SH" || true)
_has=$([ "$_count" -ge 1 ] && echo "yes" || echo "no")
assert_eq "tightened TORCH_CONSTRAINT assignment exists" "yes" "$_has"
@ -113,30 +126,47 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var"
_literal=$(grep -cE 'uv pip install .*"torch>=' "$INSTALL_SH" || true)
assert_eq "no pip install hardcodes a torch pin" "0" "$_literal"
# The same rule stated over the whole file rather than just install lines, and the two
# are complementary: the check above sees through shell-escaped quotes (\"torch>=...\")
# on an install or echoed command line, this one catches the two default ranges copied
# anywhere that is not a TORCH_CONSTRAINT= assignment, e.g. into a second variable that
# then silently diverges from the ceiling. A curated per-index override may still cap
# literally -- the gfx906 / MI50 reroute drops back below 2.11 because the rocm6.3 index
# tops out at torch 2.9.x, and the ROCm >=2.11 floors are deliberately literal -- so the
# TORCH_CONSTRAINT= assignments are the one place the literal is allowed to live.
_hardcoded=$(grep -E '"torch>=2\.4,<2\.11\.0"|"torch>=2\.4,<2\.12\.0"' "$INSTALL_SH" \
| grep -c -v '^[[:space:]]*TORCH_CONSTRAINT=' || true)
assert_eq "no hardcoded default torch range off a TORCH_CONSTRAINT= assignment" "0" "$_hardcoded"
# The gfx906 / MI50 reroute must keep its own sub-2.11 cap and must not be "fixed" to
# track the ceiling. Existence of an indented assignment, not a count: another hardware
# branch may legitimately add a second sub-2.11 cap without that being a test edit.
_count=$(grep -cE '^[[:space:]]+TORCH_CONSTRAINT="torch>=2\.4,<2\.11\.0"$' "$INSTALL_SH" || true)
_has_sub211_cap=$([ "$_count" -ge 1 ] && echo "yes" || echo "no")
assert_eq "gfx906 reroute caps torch below 2.11 for the rocm6.3 index" "yes" "$_has_sub211_cap"
# Companions must be bounded to torch's window everywhere, never bare: torchaudio 2.11
# dropped its exact torch pin, so a bare companion next to a <2.11-capped torch resolves
# a mismatched 2.11 build. Every assignment, not a fixed number of them.
# a mismatched 2.11 build. Every assignment, not a fixed number of them. The upper bound
# is a literal on the curated per-index pins and the composed ceiling on the defaults,
# so accept either form -- what is checked is that a bound is present at all.
_total=$(grep -cE '^[[:space:]]*TORCHVISION_CONSTRAINT="' "$INSTALL_SH" || true)
_bounded=$(grep -cE '^[[:space:]]*TORCHVISION_CONSTRAINT="torchvision>=[0-9][0-9.]*,<[0-9][0-9.]*"$' "$INSTALL_SH" || true)
_bounded=$(grep -cE '^[[:space:]]*TORCHVISION_CONSTRAINT="torchvision>=[0-9][0-9.]*,<([0-9][0-9.]*|[$][{]_TORCHVISION_CEILING[}])"$' "$INSTALL_SH" || true)
assert_eq "every torchvision constraint is upper-bounded" "$_total" "$_bounded"
_total=$(grep -cE '^[[:space:]]*TORCHAUDIO_CONSTRAINT="' "$INSTALL_SH" || true)
_bounded=$(grep -cE '^[[:space:]]*TORCHAUDIO_CONSTRAINT="torchaudio>=[0-9][0-9.]*,<[0-9][0-9.]*"$' "$INSTALL_SH" || true)
_bounded=$(grep -cE '^[[:space:]]*TORCHAUDIO_CONSTRAINT="torchaudio>=[0-9][0-9.]*,<([0-9][0-9.]*|[$][{]_TORCHAUDIO_CEILING[}])"$' "$INSTALL_SH" || true)
assert_eq "every torchaudio constraint is upper-bounded" "$_total" "$_bounded"
# And the top-level companion defaults specifically compose their ceiling variable
# rather than spelling a version out, so a ceiling bump stays a one-line change.
_count=$(grep -c '^TORCHVISION_CONSTRAINT="torchvision>=0.19,<${_TORCHVISION_CEILING}"$' "$INSTALL_SH" || true)
assert_eq "torchvision default composes the ceiling" "1" "$_count"
_count=$(grep -c '^TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<${_TORCHAUDIO_CEILING}"$' "$INSTALL_SH" || true)
assert_eq "torchaudio default composes the ceiling" "1" "$_count"
_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision"$' "$INSTALL_SH" || true)
assert_eq "no bare torchvision companion remains" "0" "$_count"
_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio"$' "$INSTALL_SH" || true)
assert_eq "no bare torchaudio companion remains" "0" "$_count"
# The cu* widen must carry the companions with it (torch <2.12 with torchaudio <2.11
# would cap a mismatched pair the other way).
assert_eq "cu widen pairs torchaudio (<2.12)" "1" "$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"' "$INSTALL_SH" || true)"
_gated=$(grep -c '_expected_torch_flavor_tag "$TORCH_INDEX_URL"' "$INSTALL_SH" || true)
_has_gate=$([ "$_gated" -ge 1 ] && echo "yes" || echo "no")
assert_eq "custom-companion bound gated on empty flavor tag" "yes" "$_has_gate"
# 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).
_cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true)
assert_eq "CUDA TORCH_CONSTRAINT widened to <2.12.0" "1" "$_cuda_widen"
# Widening keys off the final leaf (_torch_index_leaf), not the full URL, so a
# mirror base path with cu*/rocm7.2 but a cpu/older-rocm leaf is not mis-widened.
@ -187,7 +217,7 @@ _PS1_CONTENT=$(cat "$INSTALL_PS1")
_ps1_has_var=$(echo "$_PS1_CONTENT" | grep -c 'TORCH_CONSTRAINT\|TorchConstraint' || true)
assert_eq "install.ps1 has no TORCH_CONSTRAINT variable" "0" "$_ps1_has_var"
_ps1_hardcoded=$(echo "$_PS1_CONTENT" | grep -c '"torch>=2.4,<2.11.0"' || true)
_ps1_hardcoded=$(echo "$_PS1_CONTENT" | grep -c '"torch>=2.4,<2.12.0"' || true)
_ps1_has_hc=$([ "$_ps1_hardcoded" -ge 1 ] && echo "yes" || echo "no")
assert_eq "install.ps1 has hardcoded torch constraint" "yes" "$_ps1_has_hc"
@ -202,55 +232,55 @@ trap 'rm -rf "$TMPDIR_BASE"' EXIT
# 1. arm64 macOS py3.13 -> tightened
_result=$(run_constraint_snippet false macos arm64 13 "$TMPDIR_BASE/v1")
assert_eq "arm64+macos+py313 -> tightened" "torch>=2.6,<2.11.0" "$_result"
assert_eq "arm64+macos+py313 -> tightened" "torch>=2.6,<2.12.0" "$_result"
# 2. arm64 macOS py3.14 -> tightened (future-proofed)
_result=$(run_constraint_snippet false macos arm64 14 "$TMPDIR_BASE/v2")
assert_eq "arm64+macos+py314 -> tightened" "torch>=2.6,<2.11.0" "$_result"
assert_eq "arm64+macos+py314 -> tightened" "torch>=2.6,<2.12.0" "$_result"
# 3. arm64 macOS py3.12 -> default
_result=$(run_constraint_snippet false macos arm64 12 "$TMPDIR_BASE/v3")
assert_eq "arm64+macos+py312 -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "arm64+macos+py312 -> default" "torch>=2.4,<2.12.0" "$_result"
# 4. arm64 macOS py3.11 -> default
_result=$(run_constraint_snippet false macos arm64 11 "$TMPDIR_BASE/v4")
assert_eq "arm64+macos+py311 -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "arm64+macos+py311 -> default" "torch>=2.4,<2.12.0" "$_result"
# 5. Linux x86_64 py3.13 -> default (Linux unaffected)
_result=$(run_constraint_snippet false linux x86_64 13 "$TMPDIR_BASE/v5")
assert_eq "linux+x86_64+py313 -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "linux+x86_64+py313 -> default" "torch>=2.4,<2.12.0" "$_result"
# 6. Linux aarch64 py3.13 -> default (guard checks OS=macos)
_result=$(run_constraint_snippet false linux aarch64 13 "$TMPDIR_BASE/v6")
assert_eq "linux+aarch64+py313 -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "linux+aarch64+py313 -> default" "torch>=2.4,<2.12.0" "$_result"
# 7. Intel Mac x86_64 py3.12 -> default (arch mismatch)
_result=$(run_constraint_snippet false macos x86_64 12 "$TMPDIR_BASE/v7")
assert_eq "macos+x86_64+py312 -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "macos+x86_64+py312 -> default" "torch>=2.4,<2.12.0" "$_result"
# 8. SKIP_TORCH=true arm64 macOS py3.13 -> block skipped, default
_result=$(run_constraint_snippet true macos arm64 13 "$TMPDIR_BASE/v8")
assert_eq "SKIP_TORCH=true -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "SKIP_TORCH=true -> default" "torch>=2.4,<2.12.0" "$_result"
# 9. WSL py3.13 -> default
_result=$(run_constraint_snippet false wsl x86_64 13 "$TMPDIR_BASE/v9")
assert_eq "wsl+py313 -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "wsl+py313 -> default" "torch>=2.4,<2.12.0" "$_result"
# 10. py_minor=0 (failed query fallback) -> default
_result=$(run_constraint_snippet false macos arm64 0 "$TMPDIR_BASE/v10")
assert_eq "py_minor=0 fallback -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "py_minor=0 fallback -> default" "torch>=2.4,<2.12.0" "$_result"
# 11. Boundary: py_minor=12 -> NOT tightened
_result=$(run_constraint_snippet false macos arm64 12 "$TMPDIR_BASE/v11")
assert_eq "boundary py_minor=12 -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "boundary py_minor=12 -> default" "torch>=2.4,<2.12.0" "$_result"
# 12. Boundary: py_minor=13 -> tightened
_result=$(run_constraint_snippet false macos arm64 13 "$TMPDIR_BASE/v12")
assert_eq "boundary py_minor=13 -> tightened" "torch>=2.6,<2.11.0" "$_result"
assert_eq "boundary py_minor=13 -> tightened" "torch>=2.6,<2.12.0" "$_result"
# 13. Intel Mac py3.13 -> default (arch=x86_64, not arm64)
_result=$(run_constraint_snippet false macos x86_64 13 "$TMPDIR_BASE/v13")
assert_eq "macos+x86_64+py313 -> default" "torch>=2.4,<2.11.0" "$_result"
assert_eq "macos+x86_64+py313 -> default" "torch>=2.4,<2.12.0" "$_result"
# ======================================================================
# Mock uv integration

View file

@ -3785,12 +3785,35 @@ class TestRocmTorchPkgSpecs:
assert "2.11" in torch_spec
def test_default_caps_below_211(self):
"""Default spec (rocm7.1 and earlier) should cap below 2.11."""
"""Default spec (rocm7.0 and earlier) should cap below 2.11."""
specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("_default")
assert specs is not None
torch_spec = specs[0]
assert "<2.11" in torch_spec
def test_rocm71_repair_matches_install_sh_default_range(self):
"""rocm7.1 serves a paired 2.11 trio, so the repair path must not cap at <2.11.
install.sh leaves a rocm7.1 leaf on its default trio (torch>=2.4,<2.12.0 /
torchvision>=0.19,<0.27.0 / torchaudio>=2.4,<2.12.0), which resolves
torch 2.11.0+rocm7.1 on that index. Falling back to _default here would
force-reinstall 2.10.0+rocm7.1 over it on the next `studio update`.
"""
specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("rocm7.1")
assert specs is not None, "rocm7.1 must have its own repair spec"
assert specs == (
"torch>=2.4,<2.12.0",
"torchvision>=0.19,<0.27.0",
"torchaudio>=2.4,<2.12.0",
)
# Not the rocm7.2 spec: no 2.11 floor applies to rocm7.1.
assert specs != stack_mod._ROCM_TORCH_PKG_SPECS["rocm7.2"]
def test_rocm71_is_not_a_known_211_floor_version(self):
"""The widened rocm7.1 range must NOT promote it to a floored 2.11 line."""
assert (7, 1) not in stack_mod._ROCM_KNOWN_TORCH211_VERSIONS
assert (7, 2) in stack_mod._ROCM_KNOWN_TORCH211_VERSIONS
def test_specs_have_torch_vision_audio(self):
"""Each entry should be a 3-tuple: torch, torchvision, torchaudio."""
for tag, specs in stack_mod._ROCM_TORCH_PKG_SPECS.items():

View file

@ -0,0 +1,130 @@
#!/usr/bin/env pwsh
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Unit test for install.ps1's torch release preservation helpers
# (ConvertTo-TorchNumericRelease, Test-TorchReleaseInWindow, Get-PreviousTorchPin),
# the Windows port of install.sh's _previous_torch_pin (PR 7250). Pure helpers,
# AST-extracted and run in-process -- no GPU/venv needed.
# Run: pwsh -NoProfile -File tests/studio/test_previous_torch_pin.ps1
$ErrorActionPreference = "Stop"
$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1")
$installPath = (Resolve-Path $installPath).Path
# --- Parse install.ps1 (also serves as a syntax gate) and extract the helpers ---
$tokens = $null; $errors = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors)
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" }
foreach ($name in @("ConvertTo-TorchNumericRelease", "Test-TorchReleaseInWindow", "Get-PreviousTorchPin")) {
$fn = $ast.FindAll({ param($n)
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
}, $true)
if ($fn.Count -ne 1) { throw "expected exactly one $name in install.ps1, found $($fn.Count)" }
# Pure helpers (no exit / external calls) -- safe to define in this scope.
Invoke-Expression $fn[0].Extent.Text
}
$failures = 0
function Check($name, $cond) {
if ($cond) { Write-Host " PASS $name" }
else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ }
}
$savedUpgrade = $env:UNSLOTH_TORCH_UPGRADE
try {
Remove-Item Env:UNSLOTH_TORCH_UPGRADE -ErrorAction SilentlyContinue
$win = "torch>=2.4,<2.12.0"
# --- ConvertTo-TorchNumericRelease: accepted stable versions ---
foreach ($v in @("2.10.0", "2.10.0+cpu", "2.10.0+cu126", "2.10.1+cu130", "2.9.1+rocm7.2.1", "2.9.0+xpu", "2.10")) {
$r = ConvertTo-TorchNumericRelease $v
Check "release accepts $v" ($null -ne $r)
}
$r = ConvertTo-TorchNumericRelease "2.10.1+cu128"
Check "release strips only the +local tag" ($r.PublicBase -eq "2.10.1" -and $r.Minor -eq 10)
# --- Rejected: nightly/dev/rc/alpha/garbage never pin ---
foreach ($v in @("", " ", "2.11.0.dev20260701+cu130", "2.9.0a0+gitabc123", "2.11.0rc1+cu130",
".2.10", "2.10.", "2..10", "2.x.0", "not-a-version",
"Traceback (most recent call last):", "99999999999999999999.1")) {
Check "release rejects '$v'" ($null -eq (ConvertTo-TorchNumericRelease $v))
}
# --- Test-TorchReleaseInWindow ---
$cases = @(
@{ v = "2.4.0"; c = "torch>=2.4,<2.12.0"; ok = $true; n = "at the floor" }
@{ v = "2.3.1"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "below the floor" }
@{ v = "2.11.0"; c = "torch>=2.4,<2.12.0"; ok = $true; n = "just below the ceiling" }
@{ v = "2.12.0"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "at the ceiling" }
@{ v = "2.13.0"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "above the ceiling" }
@{ v = "2.10.0"; c = "torch>=2.11.0,<2.12.0"; ok = $false; n = "2.11 floor rejects 2.10" }
@{ v = "2.11.0"; c = "torch>=2.11.0,<2.12.0"; ok = $true; n = "2.11 floor accepts 2.11" }
@{ v = "2.11.0"; c = "torch>=2.4,<2.13.0"; ok = $true; n = "future 2.13 ceiling keeps 2.11" }
@{ v = "2.12.1"; c = "torch>=2.4,<2.13.0"; ok = $true; n = "future 2.13 ceiling keeps 2.12" }
)
foreach ($t in $cases) {
$rel = ConvertTo-TorchNumericRelease $t.v
Check ("window: " + $t.n) ((Test-TorchReleaseInWindow -Release $rel -Constraint $t.c) -eq $t.ok)
}
# Malformed constraints fail closed.
$rel = ConvertTo-TorchNumericRelease "2.10.0"
Check "window: malformed constraint fails closed" (-not (Test-TorchReleaseInWindow -Release $rel -Constraint "torch"))
Check "window: exact-pin constraint fails closed" (-not (Test-TorchReleaseInWindow -Release $rel -Constraint "torch==2.10.0"))
# --- Get-PreviousTorchPin: exact-release pin, sh parity ---
$pin = Get-PreviousTorchPin -TorchVersion "2.10.0+cu128" -Constraint $win
Check "pin keeps 2.10.0 (exact release, sh parity)" ($pin.TorchSpec -eq "torch==2.10.0")
Check "pin pairs torchvision to the kept minor" ($pin.VisionSpec -eq "torchvision==0.25.*")
Check "pin pairs torchaudio to the kept minor" ($pin.AudioSpec -eq "torchaudio==2.10.*")
$pin = Get-PreviousTorchPin -TorchVersion "2.9.1+rocm7.2.1" -Constraint $win
Check "pin keeps 2.9.1 with 0.24.*/2.9.* companions" (
$pin.TorchSpec -eq "torch==2.9.1" -and $pin.VisionSpec -eq "torchvision==0.24.*" -and $pin.AudioSpec -eq "torchaudio==2.9.*")
$pin = Get-PreviousTorchPin -TorchVersion "2.11.0+cpu" -Constraint $win
Check "pin keeps 2.11.0 under a future-widened window" ($pin.TorchSpec -eq "torch==2.11.0")
# No previous version / out-of-window / non-stable -> no pin.
Check "no pin without a previous version" ($null -eq (Get-PreviousTorchPin -TorchVersion "" -Constraint $win))
Check "no pin for a below-floor release" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.3.1+cpu" -Constraint $win))
Check "raised ROCm floor rejects keeping 2.10" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.10.0+rocm7.1" -Constraint "torch>=2.11.0,<2.12.0"))
Check "no pin for a nightly build" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.11.0.dev20260701+cu130" -Constraint $win))
Check "no pin for an unsupported 2.12 under a <2.12 window" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.12.0+cu130" -Constraint $win))
# --- UNSLOTH_TORCH_UPGRADE opt-out (exact string '1', sh parity) ---
$env:UNSLOTH_TORCH_UPGRADE = "1"
Check "UNSLOTH_TORCH_UPGRADE=1 disables the pin" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.10.0+cpu" -Constraint $win))
$env:UNSLOTH_TORCH_UPGRADE = "0"
Check "UNSLOTH_TORCH_UPGRADE=0 keeps the pin" ($null -ne (Get-PreviousTorchPin -TorchVersion "2.10.0+cpu" -Constraint $win))
} finally {
if ($null -ne $savedUpgrade) { $env:UNSLOTH_TORCH_UPGRADE = $savedUpgrade }
else { Remove-Item Env:UNSLOTH_TORCH_UPGRADE -ErrorAction SilentlyContinue }
}
# --- Structural wiring (source assertions) ---
$src = Get-Content $installPath -Raw
Check "probe runs before the rollback move" (
$src.IndexOf('$script:PrevTorchVer') -ge 0 -and
$src.IndexOf('$script:PrevTorchVer') -lt $src.IndexOf('Start-StudioVenvRollback -ExistingDir'))
Check "pin decision cites the UNSLOTH_TORCH_UPGRADE escape hatch" ($src -match 'UNSLOTH_TORCH_UPGRADE=1 to get the newest')
Check "kept-release fallback clears the pin" ($src -match '\$script:PrevTorchPin\s*=\s*\$null')
Check "kept release exported for setup.ps1" ($src -match 'UNSLOTH_KEPT_TORCH')
# The probe must read dist metadata: a broken CUDA/ROCm DLL would make "import torch"
# fail and silently drop the pin (install.sh reads metadata for the same reason).
Check "probe reads dist metadata, not import torch" (
$src -match 'importlib\.metadata as m; print\(m\.version' -and
$src -notmatch '-c "import torch; print\(torch\.__version__\)"')
# Under irm | iex the script scope is the caller's session: a second run with no existing
# venv must not inherit the earlier run's release or pin.
Check "preservation state reset before the venv branch" (
$src.IndexOf('$script:PrevTorchVer = ""') -ge 0 -and
$src.IndexOf('$script:PrevTorchVer = ""') -lt $src.IndexOf('if (Test-Path -LiteralPath $VenvPython) {'))
Check "preservation pin reset before the venv branch" (
$src.IndexOf('$script:PrevTorchPin = $null') -lt $src.IndexOf('if (Test-Path -LiteralPath $VenvPython) {'))
# A terminating exception must not leave the handoff set in a surviving session.
Check "outer finally clears the kept-torch handoff" (
$src -match '(?s)try \{\s*Install-UnslothStudio @args\s*\} finally \{.*?Remove-Item Env:UNSLOTH_KEPT_TORCH')
Write-Host ""
if ($failures -gt 0) { Write-Host "$failures check(s) failed" -ForegroundColor Red; exit 1 }
Write-Host "All checks passed"
exit 0

View file

@ -0,0 +1,134 @@
#!/usr/bin/env pwsh
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Windows twin of tests/sh/test_unsloth_torch_override.sh: install.ps1's torch-trio
# --overrides guard (New-UnslothTorchOverridesFile) on the Step-2 unsloth installs.
# The generated file also folds in the caller's UV_OVERRIDE lines, which can carry
# authenticated direct URLs, so it must never outlive the run: install.sh removes its
# twin from the EXIT/signal traps and install.ps1 must do the same from the outer
# finally. Pure text/AST assertions plus one behavioural cleanup check -- no venv needed.
# Run: pwsh -NoProfile -File tests/studio/test_unsloth_torch_override.ps1
$ErrorActionPreference = "Stop"
$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1")
$installPath = (Resolve-Path $installPath).Path
$installText = Get-Content -Raw $installPath
# --- Parse install.ps1 (also serves as a syntax gate) ---
$tokens = $null; $errors = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors)
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" }
$failures = 0
function Check($name, $cond) {
if ($cond) { Write-Host " PASS $name" }
else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ }
}
# Text of every Invoke-InstallCommandRetry statement carrying $label. A with-deps
# path has two: the overrides-guarded call and the no-torch-installed fallback.
function Get-InstallBlocks([string]$label) {
$calls = $ast.FindAll({ param($n)
$n -is [System.Management.Automation.Language.CommandAst] -and
$n.Extent.Text -like "*-Label `"$label`" *" -and
$n.GetCommandName() -eq "Invoke-InstallCommandRetry"
}, $true)
if ($calls.Count -eq 0) { throw "no Invoke-InstallCommandRetry found for '$label'" }
return @($calls | ForEach-Object { $_.Extent.Text })
}
Write-Host "New-UnslothTorchOverridesFile"
$fnAst = $ast.FindAll({ param($n)
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and
$n.Name -eq "New-UnslothTorchOverridesFile"
}, $true)
Check "helper defined exactly once" ($fnAst.Count -eq 1)
$fnText = $fnAst[0].Extent.Text
Check "helper returns null under --no-torch" ($fnText -match 'if \(\$SkipTorch\) \{ return \$null \}')
Check "helper folds in caller UV_OVERRIDE files" ($fnText -match '\$env:UV_OVERRIDE')
Write-Host "with-deps unsloth installs pass --overrides"
foreach ($label in @("install unsloth (local)", "install unsloth")) {
$blocks = Get-InstallBlocks $label
$guarded = @($blocks | Where-Object { $_ -match '--overrides \$script:TorchOverridesFile' })
Check "'$label' has one overrides-guarded call and one plain fallback" (
$blocks.Count -eq 2 -and $guarded.Count -eq 1)
}
Write-Host "the --no-deps no-torch installs carry no overrides"
foreach ($label in @("install unsloth (no-torch)", "install unsloth (migrated no-torch)")) {
$blocks = Get-InstallBlocks $label
Check "'$label' has no --overrides" (@($blocks | Where-Object { $_ -match '--overrides' }).Count -eq 0)
}
Write-Host "the generated temp file never outlives the run"
$removals = [regex]::Matches($installText, 'Remove-Item -LiteralPath \$script:TorchOverridesFile -Force')
# One in-flow removal per with-deps install, plus the outer-finally sweep.
Check "in-flow removal after each with-deps install, plus a final sweep" ($removals.Count -eq 3)
$outer = @($ast.FindAll({ param($n)
$n -is [System.Management.Automation.Language.TryStatementAst] -and
$n.Body.Extent.Text -match 'Install-UnslothStudio @args'
}, $true))
Check "outer try/finally around Install-UnslothStudio found" ($outer.Count -eq 1)
$finallyText = $outer[0].Finally.Extent.Text
Check "outer finally removes the overrides temp file" ($finallyText -match 'Remove-Item -LiteralPath \$script:TorchOverridesFile -Force')
Check "outer finally still clears UNSLOTH_KEPT_TORCH" ($finallyText -match 'Remove-Item Env:UNSLOTH_KEPT_TORCH')
# install.sh empties _UNSLOTH_TORCH_OVERRIDES before arming its traps so an inherited
# value can never be rm'd; under `irm | iex` the script scope is the caller's session,
# so the same reset must precede the outer try.
Check "overrides path reset to null before the outer try" (
$installText -match '(?m)^\$script:TorchOverridesFile = \$null\r?\ntry \{\r?\n\s*Install-UnslothStudio @args')
Write-Host "outer finally actually deletes the file after a terminating error"
# Behavioural: run the real finally body with a live temp file holding a credential-
# bearing inherited override line, exactly as an interrupted install would leave it.
$leakFile = [System.IO.Path]::GetTempFileName()
Set-Content -LiteralPath $leakFile -Encoding ascii -Value @(
"torch==2.11.0+cu128",
"private-pkg @ https://svc:TOKEN123@pkgs.corp.example/private-1.0-py3-none-any.whl")
$script:TorchOverridesFile = $leakFile
$env:UNSLOTH_KEPT_TORCH = "2.11.0"
# Strip the `finally { ... }` wrapper and run the statements themselves; Invoke-Expression
# evaluates in this scope, so the block's $script: writes land where the installer's would.
$finallyBody = ($finallyText.Trim() -replace '(?s)^\{', '') -replace '(?s)\}$', ''
try {
try { throw "simulated terminating error mid-install" }
finally { Invoke-Expression $finallyBody }
} catch { }
Check "temp overrides file removed" (-not (Test-Path -LiteralPath $leakFile))
Check "kept-torch handoff still cleared" ($null -eq $env:UNSLOTH_KEPT_TORCH)
Check "tracked path reset so a rerun cannot re-remove it" ($null -eq $script:TorchOverridesFile)
Remove-Item -LiteralPath $leakFile -Force -ErrorAction SilentlyContinue
Write-Host "the inherited-override filter drops the torch trio in any casing"
# PowerShell's -notmatch is case-insensitive unless written -cnotmatch, so a caller
# override spelled `Torch<2.11` is dropped and the generated exact pin wins.
$filterPattern = $null
if ($fnText -match '\$_ -notmatch ''([^'']+)''') { $filterPattern = $Matches[1] }
Check "filter pattern extracted from the helper" ($null -ne $filterPattern)
$inherited = @(
"# comment survives",
"Torch<2.11",
"TORCHVISION>=0.19",
"TorchAudio==2.1",
"torch<2.11.0",
"torchvision==0.25.0",
"torchaudio!=2.11.0",
"torchmetrics==1.0",
"transformers>=4.57.6",
"anyio<4.14.0"
)
$merged = @("torch==2.11.0+cu128") + @($inherited | Where-Object { $_ -notmatch $filterPattern })
$trio = @($merged | Where-Object { $_ -match '^(torch|torchvision|torchaudio)([\s<>=!~;@[]|$)' })
Check "exactly one trio requirement survives (the generated pin)" ($trio.Count -eq 1)
Check "the survivor is the generated exact pin" ($trio[0] -eq "torch==2.11.0+cu128")
foreach ($keep in @("torchmetrics==1.0", "transformers>=4.57.6", "anyio<4.14.0", "# comment survives")) {
Check "unrelated inherited override preserved: $keep" ($merged -contains $keep)
}
Write-Host ""
if ($failures -gt 0) {
Write-Host "FAILED: $failures check(s)" -ForegroundColor Red
exit 1
}
Write-Host "All checks passed."