diff --git a/install.sh b/install.sh index 92531a1f88..c52f2669a9 100755 --- a/install.sh +++ b/install.sh @@ -1832,13 +1832,45 @@ esac # Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2. case "$TORCH_INDEX_URL" in */rocm7.1|*/rocm7.1.*) - _strix_gfx="" + # Collect every gfx token in rocminfo / amd-smi enumeration order + # (skip duplicates), then index by HIP_VISIBLE_DEVICES / + # ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box + # where the user selected the dGPU does NOT get rerouted to the + # Strix per-gfx index. + _gfx_all="" if command -v rocminfo >/dev/null 2>&1; then - _strix_gfx=$(rocminfo 2>/dev/null | grep -oE 'gfx1151|gfx1150' | head -1) + _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' | awk '!seen[$0]++') fi - if [ -z "$_strix_gfx" ] && command -v amd-smi >/dev/null 2>&1; then - _strix_gfx=$(amd-smi list 2>/dev/null | grep -oE 'gfx1151|gfx1150' | head -1) + if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then + _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' | awk '!seen[$0]++') + # PowerShell paths also probe `amd-smi static --asic`; mirror it + # so a host with hipinfo-less amd-smi reports the gfx target. + if [ -z "$_gfx_all" ]; then + _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' | awk '!seen[$0]++') + fi fi + _runtime_gfx="" + if [ -n "$_gfx_all" ]; then + _vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" + _idx=0 + if [ -n "$_vis" ] && [ "$_vis" != "-1" ]; then + _first=${_vis%%,*} + case "$_first" in + ''|*[!0-9]*) _idx=0 ;; + *) _idx=$_first ;; + esac + fi + _runtime_gfx=$(printf '%s\n' "$_gfx_all" | awk -v idx="$_idx" ' + NF { vals[n++] = $0 } + END { + if (idx < 0 || idx >= n) idx = 0 + if (n > 0) print vals[idx] + }') + fi + _strix_gfx="" + case "$_runtime_gfx" in + gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;; + esac if [ -n "$_strix_gfx" ]; then echo "" >&2 echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2 diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 9a0de1cd68..59c5958389 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2060,9 +2060,24 @@ def run_training_process( global _WINDOWS_ROCM_GROUPED_MM_LIB if sys.platform == "win32": _torch_for_rocm = sys.modules.get("torch") - if _torch_for_rocm is not None and getattr( - getattr(_torch_for_rocm, "version", None), "hip", None - ): + # Broad check: torch.version.hip OR "rocm" in torch.__version__. + # AMD SDK / Radeon Windows wheels do not always populate + # torch.version.hip; without the broad check the BNB version pin, + # dynamo-disable, and _grouped_mm fallback below silently skip + # (matches the torchao stub gate above and main.py). + _build_version_for_rocm = ( + getattr(_torch_for_rocm, "__version__", "").lower() + if _torch_for_rocm is not None + else "" + ) + _is_win_rocm_torch = bool( + _torch_for_rocm is not None + and ( + getattr(getattr(_torch_for_rocm, "version", None), "hip", None) + or "rocm" in _build_version_for_rocm + ) + ) + if _is_win_rocm_torch: # Disable dynamo (belt-and-suspenders; JitDecomp patch below is the # real fix, but keeping dynamo off avoids any other compile paths). if "TORCHDYNAMO_DISABLE" not in os.environ: @@ -2112,12 +2127,24 @@ def run_training_process( # Parse HIP version for the kernel-fix gate below. # torch.version.hip can be "7.13.99004", "7.2.0", etc. - # We only need major.minor for the comparison. + # AMD SDK / Radeon wheels may leave torch.version.hip unset and + # encode the ROCm version in torch.__version__ instead + # (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back + # to that string when version.hip is missing. def _hip_ver_at_least(major: int, minor: int) -> bool: + import re as _re_ver _hip_str = getattr( getattr(_torch_for_rocm, "version", None), "hip", None ) if not _hip_str: + _ver_match = _re_ver.search( + r"rocm(\d+)\.(\d+)", _build_version_for_rocm + ) + if _ver_match: + return ( + int(_ver_match.group(1)), + int(_ver_match.group(2)), + ) >= (major, minor) return False try: _parts = [int(x) for x in str(_hip_str).split(".")[:2]] @@ -2138,11 +2165,24 @@ def run_training_process( def _grouped_mm_safe_impl( self, mat2, offs = None, bias = None, out_dtype = None ): - """Python mm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12).""" + """Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12).""" _t = _torch_for_rocm if offs is None: - # Simple case: plain matrix multiply. - result = _t.mm(self.contiguous(), mat2.contiguous()) + # No offsets: behave like the real op, which + # accepts either (M, K) x (K, N) -> mm, or 3-D + # batched inputs -> bmm. Picking torch.mm + # unconditionally previously raised "self must be + # a matrix" on 3-D MoE workloads. + if self.dim() == 3 and mat2.dim() == 3: + result = _t.bmm(self.contiguous(), mat2.contiguous()) + elif self.dim() == 3 and mat2.dim() == 2: + # Broadcast 2-D mat2 across the batch dim. + result = _t.matmul(self.contiguous(), mat2.contiguous()) + elif self.dim() == 2 and mat2.dim() == 3: + # Broadcast 2-D self across batch via matmul semantics. + result = _t.matmul(self.contiguous(), mat2.contiguous()) + else: + result = _t.mm(self.contiguous(), mat2.contiguous()) else: # Grouped case: offs[i] is the exclusive end-row of # group i in `self`; mat2 may be 3-D or 2-D. diff --git a/studio/backend/main.py b/studio/backend/main.py index 0e2175200e..fac528f1e1 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -70,25 +70,22 @@ if sys.platform == "win32": # this the server process crashes with "Configured ROCm binary not found". # Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION # before any import that pulls in bitsandbytes (mirrors worker.py logic). - # Gate on the active torch runtime, not env-var presence -- HIP_PATH / - # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a - # CUDA torch wheel, and setting BNB_ROCM_VERSION there makes bitsandbytes - # look for a ROCm DLL that doesn't exist and crash the CUDA backend. + # Gate on the active torch runtime only. AMD SDK / Radeon Windows wheels + # may not set HIP_PATH / ROCM_PATH, but they do populate torch.version.hip + # or encode "rocm" in torch.__version__. A previous version of this gate + # required HIP_PATH / ROCM_PATH and silently skipped BNB_ROCM_VERSION for + # those wheels. _is_rocm_host = False - if os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"): - try: - import torch as _torch_probe + try: + import torch as _torch_probe - # Broad check: torch.version.hip OR "rocm" in torch.__version__ -- - # AMD SDK / Radeon wheels may not populate torch.version.hip but - # still encode "rocm" in __version__. Matches worker.py + hardware.py. - _is_rocm_host = bool( - getattr(getattr(_torch_probe, "version", None), "hip", None) - or "rocm" in getattr(_torch_probe, "__version__", "").lower() - ) - del _torch_probe - except Exception: - pass + _is_rocm_host = bool( + getattr(getattr(_torch_probe, "version", None), "hip", None) + or "rocm" in getattr(_torch_probe, "__version__", "").lower() + ) + del _torch_probe + except Exception: + pass if _is_rocm_host and "BNB_ROCM_VERSION" not in os.environ: import glob as _glob diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index b9a5031374..eb2f9b763d 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -405,6 +405,26 @@ 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 / no amd-smi GUI tools), the kernel exposes AMD GPUs via + # /sys/class/kfd so `studio update` can still detect the GPU and + # repair the venv. + if sys.platform != "win32": + try: + kfd_nodes = "/sys/class/kfd/kfd/topology/nodes" + if os.path.isdir(kfd_nodes): + for entry in os.listdir(kfd_nodes): + gpu_id_path = os.path.join(kfd_nodes, entry, "gpu_id") + try: + with open(gpu_id_path) as fh: + gpu_id = fh.read().strip() + except OSError: + continue + if gpu_id and gpu_id != "0": # gpu_id 0 = CPU node + return True + except OSError: + pass return False @@ -429,30 +449,40 @@ def _has_usable_nvidia_gpu() -> bool: def _detect_amd_gfx_codes() -> list[str]: """Return the list of AMD gfx ISA strings visible to ROCm (e.g. ['gfx1151']). - Parses ``rocminfo`` output for ``ISA Info`` / ``gfx`` entries. Returns an - empty list when rocminfo is not found or no GPU agents are present. + Probes rocminfo first, then falls back to ``amd-smi list`` and + ``amd-smi static --asic`` for runtime-only Radeon hosts that ship + amd-smi but no rocminfo. Returns an empty list when no probe yields + a gfx target. """ import re - exe = shutil.which("rocminfo") - if not exe: - return [] - try: - result = subprocess.run( - [exe], - stdout = subprocess.PIPE, - stderr = subprocess.DEVNULL, - text = True, - timeout = 15, - ) - except Exception: - return [] - if result.returncode != 0: - return [] - # Match lines like " Name: gfx1151" or ISA strings - # "amdgcn-amd-amdhsa--gfx1151". Exclude the CPU agent (gfx000). - codes = re.findall(r"gfx([1-9][0-9a-z]{2,3})", result.stdout.lower()) - return list(dict.fromkeys(f"gfx{c}" for c in codes)) # deduplicate, preserve order + def _extract(text: str) -> list[str]: + codes = re.findall(r"gfx([1-9][0-9a-z]{2,3})", text.lower()) + return list(dict.fromkeys(f"gfx{c}" for c in codes)) + + probes: list[list[str]] = [] + if shutil.which("rocminfo"): + probes.append(["rocminfo"]) + if shutil.which("amd-smi"): + probes.append(["amd-smi", "list"]) + probes.append(["amd-smi", "static", "--asic"]) + for cmd in probes: + try: + result = subprocess.run( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 15, + ) + except Exception: + continue + if result.returncode != 0 or not result.stdout.strip(): + continue + codes = _extract(result.stdout) + if codes: + return codes + return [] # Set by _ensure_rocm_torch() on success; suppresses the post-install AMD warning. @@ -709,13 +739,49 @@ def _ensure_rocm_torch() -> None: f" skipping AMD per-gfx index override.\n" ) - if not has_hip_torch: - 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}") + # Strix override on ROCm 7.1 must fire even when has_hip_torch is True -- + # an existing torch with `torch.version.hip == "7.1"` is exactly the broken + # combo the override is meant to repair, so skipping it leaves users on + # the known _grouped_mm segfault. + 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( + "ROCm torch (Strix arch-specific)", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + index_url, + constrain = False, + ) + rocm_torch_ready = True + elif 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"] + ) pip_install( - "ROCm torch (Strix arch-specific)", + f"ROCm torch ({tag})", "--force-reinstall", "--no-cache-dir", _torch_pkg, @@ -726,39 +792,6 @@ 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/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 3b3ebd206d..6e759afecc 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2338,7 +2338,7 @@ class TestStrixRocm71Override: strix_idx = source.find("_strix_gfx") assert strix_idx != -1 # Look back for the rocm7.1 pattern within 600 chars before _strix_gfx - context_before = source[max(0, strix_idx - 600) : strix_idx] + context_before = source[max(0, strix_idx - 2400) : strix_idx] assert "rocm7.1" in context_before def test_torch_constraint_updated_for_strix_amd_index(self):