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

fix(install): route Strix to AMD gfx index on ROCm 7.14 (#7300)

Browse source 
* fix(install): route Strix to AMD gfx index on ROCm 7.14

When ROCm 7.3+ caps to the generic pytorch.org rocm7.2 index (or the
Radeon repo is unavailable), gfx1150/gfx1151 hosts were left on
torch 2.11+rocm7.2 instead of AMD's arch-specific wheels. Broaden the
Strix reroute in install.sh and studio/install_python_stack.py so
`studio update` repairs the same path as fresh installs (unslothai#7280).

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Souravrajvi0 2026-07-22 17:33:51 +05:30 • committed by GitHub
commit 84b762228c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 59 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

37
studio/install_python_stack.py
View file

@ -73,6 +73,27 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
(6, 0): "rocm6.0",
}
def _generic_pytorch_rocm_tag(ver: tuple[int, int]) -> str | None:
"""Newest download.pytorch.org rocmX.Y tag for a host ROCm version."""
return next(
(t for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) if ver >= (maj, mn)),
None,
)
_ROCM_ARCH_INDEX_FLOOR = (7, 13) # AMD per-arch index ships torch 2.11+rocm7.13
def _strix_needs_amd_arch_index(ver: tuple[int, int]) -> bool:
"""True when Strix's generic pytorch.org index sits below the AMD arch floor
(7.13), so gfx1150/1151 must use repo.amd.com's per-arch wheels. Mirrors
install.sh _rocm_leaf_below: reroute any generic rocm index (6.x/7.0/7.2 and a
future 7.3+), never one at/above the floor."""
key = next((k for k in sorted(_ROCM_TORCH_INDEX, reverse = True) if ver >= k), None)
return key is not None and key < _ROCM_ARCH_INDEX_FLOOR
# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug).
# Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare.
_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"})
@ -1691,13 +1712,13 @@ def _ensure_rocm_torch() -> None:
rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch
# Strix Halo / Point (gfx1151 / gfx1150) segfault under ROCm 7.1 in torch._grouped_mm;
# AMD's per-gfx repo ships 2.11.0+rocm7.13.0 with the fix, so route those hosts there
# (mirrors install.sh). On mixed hosts, reroute only when HIP's runtime GPU is the Strix one.
# Strix Halo / Point (gfx1151 / gfx1150) need torch from AMD's per-gfx index
# (2.11+rocm7.13); any generic pytorch.org rocm index lacks the fixes (ROCm 7.1
# segfaults in _grouped_mm). See _strix_needs_amd_arch_index for the floor gate.
_strix_override_url: "str | None" = None
_strix_override_pkgs: "tuple[str, str, str] | None" = None
# An explicit ROCm pin is authoritative: never auto-reroute it.
if ver < (7, 2) and _explicit_rocm_torch_index_url() is None:
if _strix_needs_amd_arch_index(ver) and _explicit_rocm_torch_index_url() is None:
gfx_codes = _detect_amd_gfx_codes()
_strix_gfx = {"gfx1151", "gfx1150"}
_detected_strix = _strix_gfx.intersection(gfx_codes)
@ -1721,10 +1742,10 @@ def _ensure_rocm_torch() -> None:
print(
f"\n {_selected_gfx} (AMD Strix) is the runtime target with ROCm "
f"{ver[0]}.{ver[1]}.\n"
f" ROCm 7.1 has a known _grouped_mm segfault on this GPU;\n"
f" routing torch install to AMD's arch-specific index\n"
f" Routing torch install to AMD's arch-specific index\n"
f" ({_strix_override_url}) which serves torch 2.11.0+rocm7.13.0\n"
f" with the upstream fix.\n"
f" with AMD's gfx1150/gfx1151 fixes (more reliable than the generic\n"
f" pytorch.org rocm7.2 index on ROCm 7.3+ hosts).\n"
)
else:
_gfx_str = ", ".join(sorted(_detected_strix))
@ -1740,7 +1761,7 @@ def _ensure_rocm_torch() -> None:
index_url = _strix_override_url
_torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs
print(
f" Strix ROCm 7.1 override -- installing torch from "
f" Strix arch-specific override -- installing torch from "
f"{_strip_index_url_credentials(index_url)}"
)
pip_install(

30
tests/studio/install/test_rocm_support.py
View file

@ -710,6 +710,27 @@ class TestEnsureRocmTorch:
torch_call = mock_pip.call_args_list[0]
assert "rocm7.2" in str(torch_call)
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 14))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1150"])
def test_rocm_714_strix_routes_to_amd_arch_index(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""ROCm 7.14 caps to rocm7.2 on pytorch.org; Strix must use AMD gfx index."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"7.14.60850|2.11.0+rocm7.2\n"
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "gfx1150" in torch_call
assert "torch>=2.11.0,<2.12.0" in torch_call
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@ -3253,6 +3274,15 @@ class TestStrixRocm71Override:
assert r.returncode == 0, f"probe aborted under set -e: {r.stderr}"
assert "OK:gfx1151" in r.stdout, f"amd-smi fallback not reached: {r.stdout!r}"
def test_strix_routing_helpers_cover_rocm714(self):
# Reroute for any generic pytorch.org index below the 7.13 arch floor (7.0,
# 7.2, a future 7.3+), never at/above it -- mirrors install.sh _rocm_leaf_below.
assert stack_mod._generic_pytorch_rocm_tag((7, 14)) == "rocm7.2"
assert stack_mod._strix_needs_amd_arch_index((7, 14)) is True
assert stack_mod._strix_needs_amd_arch_index((7, 0)) is True
assert stack_mod._strix_needs_amd_arch_index((6, 0)) is True
assert stack_mod._strix_needs_amd_arch_index((5, 0)) is False
def test_torch_constraint_updated_for_strix_amd_index(self):
"""install.sh must set TORCH_CONSTRAINT>=2.11 when routing Strix to AMD index."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")