diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index b58e94cd3f..bb329e189e 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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( diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 5825bbe31f..b343b07238 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -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")