From 21773215d93ce107ae9b01dd40c2fb0cfb2e915b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 19 May 2026 10:22:23 +0000 Subject: [PATCH] fix(studio/rocm): robustness pass - rocm tag normalisation, Strix routing parity, hardened detection Robustness pass on top of 76137b2d. Four targeted fixes: 1. install.sh ROCm-tag routing normalisation. `rocm7.2.1` would route to https://download.pytorch.org/whl/rocm7.2.1 which does not exist (PyTorch publishes major.minor URLs only). Same for any future patch-level tag. Normalise every rocm{maj.min}* pattern to the bare {maj.min} index URL. 2. install.ps1 + studio/setup.ps1 marketing-name fallback. The gfx1151 row matched 890M / Strix Halo / HX 37x / HX 38x / AI 9 HX but not the actual retail name 'AMD Radeon 8060S Graphics' shipped by OEMs (Ryzen AI MAX+ 395). Add '8060S' to the regex. 3. install_python_stack.py Strix + ROCm 7.1 routing parity with install.sh. The shell installer reroutes Strix Halo / Point + ROCm 7.1 to repo.amd.com/rocm/whl/{gfx}/ (which serves torch 2.11.0+rocm7.13.0 with the upstream _grouped_mm fix). The Python `studio update` path only warned and still installed the broken generic rocm7.1 wheel. Mirror the override: detect gfx1151/gfx1150 on ROCm 7.1, route to the AMD per-gfx index, honour UNSLOTH_AMD_ROCM_MIRROR override. 4. _detect_windows_gfx_arch amd-smi parsing tightened. The amd-smi fallback added in the prior commit used a bare `\bgfx[1-9][0-9a-z]{2,3}\b` match against the lowercased stdout, which could pick up stray gfx references in warnings / device-name strings. Anchor on labelled lines first (Target_Graphics_Version, ASIC, Arch, gfx) and fall back to the bare match only when no labelled line is present. Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py; sim_5301 23 cases pass (6 new sims for the Strix override + amd-smi parsing). --- install.ps1 | 2 +- install.sh | 12 +++- studio/install_python_stack.py | 110 +++++++++++++++++++++++---------- studio/setup.ps1 | 2 +- 4 files changed, 90 insertions(+), 36 deletions(-) diff --git a/install.ps1 b/install.ps1 index 361813f862..a3309485f3 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1307,7 +1307,7 @@ shell.Run cmd, 0, False $nameArchTable = @( @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 - @{ P = "890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo) + @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop @{ P = "RX 7600"; A = "gfx1102" } # RDNA 3 diff --git a/install.sh b/install.sh index 72a702dc37..a13ed0a3d0 100755 --- a/install.sh +++ b/install.sh @@ -1644,9 +1644,17 @@ get_torch_index_url() { echo "$_base/cpu"; return ;; esac # Supported tags; 6.5+ clips to rocm6.4, 7.3+ caps to rocm7.2. + # PyTorch publishes major.minor URLs only (no patch level), so + # rocm7.2.1 / rocm6.0.2 / etc. must normalise to rocm7.2 / rocm6.0. case "$_rocm_tag" in - rocm6.0|rocm6.0.*|rocm6.1|rocm6.1.*|rocm6.2|rocm6.2.*|rocm6.3|rocm6.3.*|rocm6.4|rocm6.4.*|rocm7.0|rocm7.0.*|rocm7.1|rocm7.1.*|rocm7.2|rocm7.2.*) - echo "$_base/$_rocm_tag" ;; + rocm6.0*) echo "$_base/rocm6.0" ;; + rocm6.1*) echo "$_base/rocm6.1" ;; + rocm6.2*) echo "$_base/rocm6.2" ;; + rocm6.3*) echo "$_base/rocm6.3" ;; + rocm6.4*) echo "$_base/rocm6.4" ;; + rocm7.0*) echo "$_base/rocm7.0" ;; + rocm7.1*) echo "$_base/rocm7.1" ;; + rocm7.2*) echo "$_base/rocm7.2" ;; rocm6.*) # ROCm 6.5+ (no published PyTorch wheels): clip down # to the last supported 6.x wheel set. diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 139c496793..c2bdd30370 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -282,10 +282,21 @@ def _detect_windows_gfx_arch() -> str | None: ) if result.returncode != 0: continue - text = result.stdout.decode(errors = "replace").lower() - m = re.search(r"\bgfx[1-9][0-9a-z]{2,3}\b", text) - if m: - return m.group(0) + text = result.stdout.decode(errors = "replace") + # Anchor on a labelled gfx line (e.g. "TARGET_GRAPHICS_VERSION: gfx1151" + # or "Arch: gfx1151") to avoid catching stray gfx mentions in + # warnings or device-name strings. Fall back to a bare token match + # only if no labelled line is found. + m = re.search( + r"(?im)^\s*(?:target_graphics_version|gfx|arch|asic)\b[^:\r\n]*:\s*(gfx[1-9][0-9a-z]{2,3})\b", + text, + ) + if not m: + m = re.search(r"\bgfx[1-9][0-9a-z]{2,3}\b", text.lower()) + if m: + return m.group(0) + continue + return m.group(1).lower() except Exception: continue return None @@ -610,42 +621,44 @@ def _ensure_rocm_torch() -> None: rocm_torch_ready = has_hip_torch - # Strix Halo (gfx1151) segfaults under ROCm 7.1 due to a ROCm driver bug - # fixed in ROCm 7.2. Warn early so users know why training may crash. + # Strix Halo / Strix Point (gfx1151 / gfx1150) segfault under ROCm 7.1 + # in torch._grouped_mm. AMD's per-gfx repo ships torch 2.11.0+rocm7.13.0 + # with the real fix, so route those hosts there instead of the generic + # pytorch.org rocm7.1 wheel. Mirrors install.sh's Strix override. + _strix_override_url: "str | None" = None + _strix_override_pkgs: "tuple[str, str, str] | None" = None if ver < (7, 2): gfx_codes = _detect_amd_gfx_codes() _strix_gfx = {"gfx1151", "gfx1150"} - if _strix_gfx.intersection(gfx_codes): - _gfx_str = ", ".join(sorted(_strix_gfx.intersection(gfx_codes))) + _detected_strix = _strix_gfx.intersection(gfx_codes) + if _detected_strix: + _gfx_str = ", ".join(sorted(_detected_strix)) + _selected_gfx = sorted(_detected_strix)[0] + _amd_mirror = ( + os.environ.get("UNSLOTH_AMD_ROCM_MIRROR") + or "https://repo.amd.com/rocm/whl" + ).rstrip("/") + _strix_override_url = f"{_amd_mirror}/{_selected_gfx}/" + _strix_override_pkgs = ( + "torch>=2.11.0,<2.12.0", + "torchvision", + "torchaudio", + ) print( - f"\n ⚠️ {_gfx_str} (AMD Strix Halo) detected with ROCm {ver[0]}.{ver[1]}.\n" - f" ROCm 7.1 has a known segfault on this GPU when tensors are\n" - f" moved to the GPU. Upgrade to ROCm 7.2+ to enable training.\n" + f"\n ⚠️ {_gfx_str} (AMD Strix) detected with ROCm {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" ({_strix_override_url}) which serves torch 2.11.0+rocm7.13.0\n" + f" with the upstream fix.\n" ) if not has_hip_torch: - # Select best matching wheel tag (newest ROCm version <= installed) - tag = next( - ( - t - for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) - if ver >= (maj, mn) - ), - None, - ) - if tag is None: - print( - f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " - f"skipping torch reinstall" - ) - else: - index_url = f"{_PYTORCH_WHL_BASE}/{tag}" - print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}") - _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( - tag, _ROCM_TORCH_PKG_SPECS["_default"] - ) + 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 + print(f" Strix ROCm 7.1 override -- installing torch from {index_url}") pip_install( - f"ROCm torch ({tag})", + "ROCm torch (Strix arch-specific)", "--force-reinstall", "--no-cache-dir", _torch_pkg, @@ -656,6 +669,39 @@ def _ensure_rocm_torch() -> None: constrain = False, ) rocm_torch_ready = True + else: + # Select best matching wheel tag (newest ROCm version <= installed) + tag = next( + ( + t + for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) + if ver >= (maj, mn) + ), + None, + ) + if tag is None: + print( + f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " + f"skipping torch reinstall" + ) + else: + index_url = f"{_PYTORCH_WHL_BASE}/{tag}" + print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}") + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( + tag, _ROCM_TORCH_PKG_SPECS["_default"] + ) + pip_install( + f"ROCm torch ({tag})", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + index_url, + constrain = False, + ) + rocm_torch_ready = True # Install bitsandbytes only when torch links against ROCm. Prefers the # continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 2772c94b11..6ee20143ec 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -792,7 +792,7 @@ if (-not $HasNvidiaSmi) { $nameArchTable = @( @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 - @{ P = "890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo) + @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop @{ P = "RX 7600"; A = "gfx1102" } # RDNA 3