fix(studio/rocm): worker BNB/grouped_mm broad gate, install.sh Strix visibility, runtime-only ROCm detection
Round-5 robustness pass based on 20 parallel reviewers of head 96b9e465.
1. studio/backend/core/training/worker.py - BNB version pin / dynamo disable
/ _grouped_mm fallback block was still gated on torch.version.hip alone
despite the torchao stub block above already using the broad check. AMD
SDK / Radeon Windows wheels (torch.__version__ contains "rocm" but
torch.version.hip is None) silently skipped the Windows ROCm runtime
patches. Aligned to the same broad check (8/20 reviewers).
2. studio/backend/core/training/worker.py - _hip_ver_at_least() now also
parses the ROCm version out of torch.__version__ (e.g. "2.11.0+rocm7.13.0")
when torch.version.hip is missing, so the kernel-fix gate is correct for
SDK / Radeon wheels too.
3. studio/backend/core/training/worker.py - _grouped_mm_safe_impl with
offs=None now picks torch.bmm/matmul for 3-D inputs instead of always
calling torch.mm. The real _grouped_mm accepts 3-D batched matmul; the
prior fallback raised "self must be a matrix" on MoE workloads (2/20).
4. studio/backend/main.py - dropped the HIP_PATH / ROCM_PATH env-var gate
from the BNB block; probe torch directly. Runtime-only Radeon / AMD SDK
Windows installs do not set those SDK env vars but still ship ROCm torch
(5/20 reviewers).
5. install.sh - Strix override now collects every gfx token from
rocminfo / amd-smi (in enumeration order), then indexes by
HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-
Strix dGPU host where the user selected the dGPU does NOT get rerouted
to the Strix per-gfx index. Mirrors the Python update path (5/20 reviewers).
6. install.sh - Strix detection chain now also probes `amd-smi static --asic`,
matching the PowerShell installer (1/20). Closes the gap on runtime-only
Strix hosts where `amd-smi list` does not surface a gfx token.
7. studio/install_python_stack.py - _has_rocm_gpu() now has the sysfs KFD
topology fallback (/sys/class/kfd/kfd/topology/nodes/*/gpu_id), matching
install.sh. On minimal package-managed installs without rocminfo /
amd-smi GUI tools, `studio update` can now detect the GPU and repair the
venv instead of returning early (2/20).
8. studio/install_python_stack.py - _detect_amd_gfx_codes() now falls back
to `amd-smi list` and `amd-smi static --asic` when rocminfo is missing
(2/20). Strix routing on runtime-only Radeon hosts now matches what
install.sh has done for a while.
9. studio/install_python_stack.py - Strix override now applies even when
has_hip_torch is True. The whole point of the override is to repair an
existing broken torch.version.hip == "7.1" install; skipping the
reinstall left users on the known _grouped_mm segfaulting stack (3/20).
Tests: 231 passed, 1 skipped. sim_5301 30 cases pass. sim_cross 12 pass.
This commit is contained in:
Daniel Han
parent
825cbf5f21
commit
8c3024133d
5 changed files with 192 additions and 90 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue