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

install_python_stack.py: keep the ROCm 7.1 repair on the 2.11 line

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-26 15:50:14 +00:00
commit 1e3b1c97ad
2 changed files with 40 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

20
studio/install_python_stack.py
View file

@ -58,8 +58,8 @@ PLATFORM_LACKS_TORCHCODEC_WHEEL = (
# 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",
@ -97,14 +97,26 @@ _ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset(
# 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",