fix(install): infer Strix gfx when ROCm runtime is absent (#7305)

* fix(install): infer Strix gfx when ROCm runtime is absent

When /dev/kfd and rocminfo are missing on Linux (e.g. Arch/CachyOS Strix
Halo), route to AMD per-arch wheels via cpuinfo/lspci inference instead
of CPU-only PyTorch. Mirrors install.ps1 Windows behavior and fixes
studio update via install_python_stack.py (unslothai#7301).

* Map Radeon 8065S to gfx1151 in the Linux gfx inference (Codex P2)

install.sh _infer_amd_gfx_arch_from_gpu_name missed 8065S, so a Strix Halo host that only exposes 'AMD Radeon 8065S' via lspci (no Ryzen AI Max branding in /proc/cpuinfo) was left on CPU torch. setup.sh and setup.ps1 already list 8065S -> gfx1151. Added it, and widened the cpuinfo regexes (install.sh and install_python_stack.py) from Radeon 80[0-9]0S to 80[0-9][05]S to match the 80X5S naming, consistent with the display-side check already in install.sh. Tests cover the 8065S name and the cpuinfo-only case.

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

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

* Gate the Linux gfx inference out of WSL without the ROCDXG runtime for PR #7305

On WSL /proc/cpuinfo and lspci still see the host APU, so a standalone
'unsloth studio update' could infer gfx1151 and install per-arch ROCm wheels
into a WSL env whose ROCDXG bridge (librocdxg) was never bootstrapped, i.e. one
that cannot expose the GPU. Skip the cpuinfo/lspci inference on WSL unless
librocdxg is present; an explicit UNSLOTH_ROCM_GFX_ARCH override still wins.

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

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

* Address Codex review on PR #7305 (WSL runtime gate, Linux mirror, arch guard)

- install.sh _infer_linux_amd_gfx_arch: skip the cpuinfo/lspci inference on WSL
  unless librocdxg is present (the ROCDXG bridge), mirroring the Python fix, so a
  WSL box whose ROCm bootstrap was skipped keeps the CPU fallback instead of
  installing AMD wheels that cannot reach the GPU. The explicit UNSLOTH_ROCM_GFX_ARCH
  override still returns first, so it stays authoritative.
- install.sh: guard the inferred-gfx reroute on x86_64|amd64. ROCm torch wheels are
  not published for arm64, so an inferred/overridden gfx no longer pushes an arm64
  host to the AMD arch index (get_torch_index_url returns CPU there).
- install_python_stack.py _amd_arch_index_url: honour UNSLOTH_AMD_ROCM_MIRROR on
  Linux (the same var install.sh uses) instead of the Windows mirror var, so a
  mirrored/air-gapped Linux 'unsloth studio update' reaches the index install.sh
  chose. Windows still delegates unchanged; both default to repo.amd.com.

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

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

* Scan all AMD display controllers in the lspci fallback for PR #7305 (Codex P2)

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

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

* fix(studio): keep inferred AMD wheels from being overwritten

After a successful inferred-gfx install, skip the generic pytorch.org
ROCm reinstall so readable ROCm userland without /dev/kfd cannot undo
the per-arch repair (Codex P1 on #7305). Also merge latest main.

* Only take the inferred-gfx install when the runtime sees no GPU for PR #7305 (Codex P1)

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

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

* Isolate three updater tests from the host cpuinfo for PR #7305 (Strix dev box leak)

* Gate the reroute on invisible ROCm and forward the inferred gfx to setup.sh for PR #7305 (Codex P2s)

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

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

* Require AMD PCI display evidence for cpuinfo inference; honor gfx override with visible ROCm for PR #7305 (Codex P2s)

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: LeoBorcherding <borchborchmail@gmail.com>
This commit is contained in:
Souravrajvi0 2026-07-23 06:46:45 +05:30 committed by GitHub
commit 978ae4745b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 714 additions and 8 deletions

View file

@ -2144,6 +2144,92 @@ _amd_gpu_present_via_pci() {
return 1
}
# Map a gfx arch to the AMD pip index family (mirrors install.ps1 $archFamilyMap).
_amd_arch_index_family_for_gfx() {
case "$1" in
gfx1201|gfx1200) echo gfx120X-all ;;
gfx1151) echo gfx1151 ;;
gfx1150) echo gfx1150 ;;
gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;;
gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;;
gfx90a) echo gfx90a ;;
gfx908) echo gfx908 ;;
*) return 1 ;;
esac
}
# Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable).
_infer_amd_gfx_arch_from_gpu_name() {
case "$1" in
*"9070 XT"*|*9080*) echo gfx1201 ;;
*9070*|*9060*) echo gfx1200 ;;
*"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;;
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1150 ;;
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) echo gfx1102 ;;
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) echo gfx1100 ;;
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;;
*"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;;
*"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;;
*"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) echo gfx1034 ;;
*) return 1 ;;
esac
}
# Best-effort gfx inference when ROCm tools can't see the GPU (unslothai#7301).
# Mirrors install.ps1 arch resolution on Windows ($HasROCm false, $ROCmGfxArch set).
_infer_linux_amd_gfx_arch() {
if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then
printf '%s\n' "$(printf '%s' "$UNSLOTH_ROCM_GFX_ARCH" | tr '[:upper:]' '[:lower:]')"
return 0
fi
# On WSL /proc/cpuinfo and lspci still report the host APU, but without the
# ROCDXG bridge (librocdxg over /dev/dxg) the AMD wheels can't reach the GPU;
# keep the CPU fallback there unless that runtime is present (the explicit
# override above still wins). Mirrors install_python_stack.py.
_gpu_evidence=""
if [ -e /dev/dxg ] || grep -qi microsoft /proc/version 2>/dev/null; then
for _d in /opt/rocm/lib /opt/rocm/lib64 /opt/rocm-*/lib /opt/rocm-*/lib64; do
{ [ -e "$_d/librocdxg.so" ] || [ -e "$_d/librocdxg.so.1" ]; } && _rocdxg=1 && break
done
[ -n "${_rocdxg:-}" ] || return 1
# WSL enumerates no PCI display device; /dev/dxg + librocdxg IS the
# GPU evidence there.
_gpu_evidence=1
elif _amd_gpu_present_via_pci; then
_gpu_evidence=1
fi
# /proc/cpuinfo leaks the HOST CPU model into VMs/containers that received
# no AMD GPU, so the CPU-model text alone is not GPU evidence: require an
# AMD display device (PCI vendor 0x1002, class 0x03*) before trusting it.
# The lspci fallback below needs no gate; an AMD display line IS evidence.
if [ -n "$_gpu_evidence" ] && grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null; then
echo gfx1151
return 0
fi
if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then
echo gfx1150
return 0
fi
if command -v lspci >/dev/null 2>&1; then
# A non-AMD controller can enumerate first (Intel/ASPEED before an AMD
# dGPU), so scan every display-class line and take the first AMD one
# that maps. The vendor guard is case-SENSITIVE (a -i "ATI" would match
# "CorporATIon" on every Intel/NVIDIA line); whole-line matching also
# survives the 0000: PCI domain prefix. Mirrors install_python_stack.py.
_amd_disp=$(lspci -nn 2>/dev/null | grep -E 'VGA compatible controller|3D controller|Display controller' | grep -E 'AMD|ATI' || true)
while IFS= read -r _ln; do
[ -n "$_ln" ] || continue
if _gfx=$(_infer_amd_gfx_arch_from_gpu_name "$_ln"); then
echo "$_gfx"
return 0
fi
done <<EOF
$_amd_disp
EOF
fi
return 1
}
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@ -2752,6 +2838,60 @@ fi
TORCH_INDEX_URL=$(get_torch_index_url)
# Linux: ROCm runtime missing but a supported AMD gfx arch is inferable (Strix Halo
# in /proc/cpuinfo, lspci marketing name, UNSLOTH_ROCM_GFX_ARCH). Route to AMD's
# per-arch wheels like install.ps1 does on Windows (unslothai#7301).
# Gated on _has_amd_rocm_gpu being FALSE: a */cpu index on a host whose GPU IS
# visible to the ROCm probes is a deliberate fallback (unsupported/unreadable
# ROCm version, after its own warning), not a missing runtime -- rerouting it
# would contradict that decision. An explicit UNSLOTH_ROCM_GFX_ARCH override
# stays authoritative either way.
if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \
! _has_usable_nvidia_gpu && \
{ [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu; } && \
case "$(uname -s)" in Linux) true ;; *) false ;; esac && \
case "$_ARCH" in x86_64|amd64) true ;; *) false ;; esac; then
# ROCm torch wheels are x86_64-only; get_torch_index_url returns CPU on other
# arches, so an inferred/overridden gfx must not reroute arm64 to AMD wheels.
case "$TORCH_INDEX_URL" in
*/cpu)
_linux_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null || true)
if [ -n "$_linux_inferred_gfx" ]; then
_amd_family=$(_amd_arch_index_family_for_gfx "$_linux_inferred_gfx") || _amd_family=""
if [ -n "$_amd_family" ]; then
_amd_mirror="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}"
while [ "${_amd_mirror%/}" != "$_amd_mirror" ]; do
_amd_mirror="${_amd_mirror%/}"
done
TORCH_INDEX_URL="${_amd_mirror}/${_amd_family}/"
# Hand the inferred arch to setup.sh (llama.cpp): it re-probes
# ROCm on its own, and on these runtime-less hosts its probes
# find nothing, so without this it classifies the box as
# non-ROCm and installs the CPU prebuilt while torch just got
# AMD per-arch wheels. setup.sh and install_llama_prebuilt.py
# both honor UNSLOTH_ROCM_GFX_ARCH, so exporting it is the
# whole handoff (a user-set override re-exports unchanged).
export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx"
case "$_linux_inferred_gfx" in
gfx1201|gfx1200|gfx1151|gfx1150)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
;;
esac
echo "" >&2
echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2
echo " [WARN] Routing to AMD arch-specific wheels ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))." >&2
echo " [WARN] These wheels bundle their own ROCm runtime; install the kernel stack for native compute:" >&2
echo " [WARN] https://docs.unsloth.ai/get-started/install-and-update/amd" >&2
echo " [WARN] Tip: set UNSLOTH_ROCM_GFX_ARCH=$_linux_inferred_gfx to skip inference next time." >&2
echo "" >&2
fi
fi
;;
esac
fi
# Export the resolved torch backend ("cuda", "rocm", or "cpu") so that
# downstream scripts (setup.sh -> install_python_stack.py) know what was
# chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts.

View file

@ -769,6 +769,142 @@ def _gfx_arch_from_gpu_name(name: str) -> "str | None":
return None
def _linux_amd_gfx_from_cpuinfo() -> "str | None":
"""Infer gfx arch from /proc/cpuinfo on integrated AMD APUs (Strix Halo/Point)."""
try:
text = Path("/proc/cpuinfo").read_text(encoding = "utf-8", errors = "replace")
except OSError:
return None
if re.search(r"Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo", text, re.IGNORECASE):
return "gfx1151"
if re.search(
r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]"
r"|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33",
text,
re.IGNORECASE,
):
return "gfx1150"
return None
def _linux_amd_gfx_from_lspci() -> "str | None":
"""First AMD display-class lspci line mapping to a known gfx arch. A non-AMD
controller can enumerate first (Intel/ASPEED before an AMD dGPU), so scan
them all. The vendor guard is case-SENSITIVE: a -i "ATI" would match
"CorporATIon" on every Intel/NVIDIA line. Whole-line matching also survives
the 0000: PCI domain prefix."""
lspci = shutil.which("lspci")
if not lspci:
return None
try:
result = subprocess.run(
[lspci, "-nn"],
stdout = subprocess.PIPE,
stderr = subprocess.DEVNULL,
text = True,
timeout = 10,
)
except Exception:
return None
if result.returncode != 0:
return None
for line in result.stdout.splitlines():
if not re.search(r"VGA compatible controller|3D controller|Display controller", line, re.I):
continue
if not re.search(r"AMD|ATI", line):
continue
arch = _gfx_arch_from_gpu_name(line)
if arch:
return arch
return None
def _is_wsl() -> bool:
"""True on WSL, where the AMD GPU is reached via /dev/dxg (not /dev/kfd)."""
if os.path.exists("/dev/dxg"):
return True
try:
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
return "microsoft" in fh.read().lower()
except OSError:
return False
def _wsl_rocm_runtime_present() -> bool:
"""librocdxg (the WSL ROCDXG bridge that lets HIP reach the GPU over /dev/dxg)
under a ROCm lib dir. Its absence marks a WSL box whose ROCm was never set up."""
dirs = ["/opt/rocm/lib", "/opt/rocm/lib64"]
dirs += glob.glob("/opt/rocm-*/lib") + glob.glob("/opt/rocm-*/lib64")
return any(
os.path.exists(os.path.join(d, so))
for d in dirs
for so in ("librocdxg.so", "librocdxg.so.1")
)
def _linux_amd_display_device_present() -> bool:
"""Any AMD (vendor 0x1002) PCI display-class (0x03*) device in sysfs.
/proc/cpuinfo leaks the HOST CPU model into VMs/containers that received no
AMD GPU, so the CPU-model text alone is not GPU evidence; this is the
device-level check (mirrors install.sh _amd_gpu_present_via_pci)."""
try:
for dev in Path("/sys/bus/pci/devices").iterdir():
try:
if (dev / "vendor").read_text().strip() != "0x1002":
continue
if (dev / "class").read_text().strip().startswith("0x03"):
return True
except OSError:
continue
except OSError:
pass
return False
def _infer_linux_amd_gfx_arch() -> "str | None":
"""Infer gfx when ROCm runtime is absent but the host is a known AMD arch (unslothai#7301)."""
override = (os.environ.get("UNSLOTH_ROCM_GFX_ARCH") or "").strip().lower()
if override:
return override
if _is_wsl():
# cpuinfo/lspci see the host APU even on a WSL box whose ROCDXG runtime
# was never bootstrapped; inferring there would install per-arch ROCm
# wheels into an env that still can't expose the GPU. Skip unless that
# runtime is present -- WSL enumerates no PCI display device, so
# /dev/dxg + librocdxg IS the GPU evidence there.
if not _wsl_rocm_runtime_present():
return None
elif not _linux_amd_display_device_present():
# Native Linux: a VM/container on a Strix host still shows the host CPU
# model in /proc/cpuinfo while receiving no AMD GPU, so require an AMD
# display device before trusting the CPU-model inference. The lspci
# fallback reads the same PCI space and would find nothing here either.
return None
cpu_gfx = _linux_amd_gfx_from_cpuinfo()
if cpu_gfx:
return cpu_gfx
return _linux_amd_gfx_from_lspci()
def _amd_arch_index_url(gfx_arch: str | None) -> str | None:
"""Return the AMD per-arch pip index URL for a gfx arch (Linux + Windows).
Windows honors UNSLOTH_ROCM_WINDOWS_MIRROR (via _windows_rocm_index_url);
Linux honors UNSLOTH_AMD_ROCM_MIRROR -- the same var install.sh uses -- so a
mirrored/air-gapped Linux repair reaches the index install.sh chose rather
than falling back to repo.amd.com. Both default to repo.amd.com when unset.
"""
if IS_WINDOWS:
return _windows_rocm_index_url(gfx_arch)
arch_family = _GFX_TO_AMD_INDEX_ARCH.get(gfx_arch or "")
if arch_family is None:
return None
base = (os.environ.get("UNSLOTH_AMD_ROCM_MIRROR") or "https://repo.amd.com/rocm/whl").rstrip(
"/"
)
return f"{base}/{arch_family}/"
def _windows_rocm_index_url(gfx_arch: str | None) -> str | None:
"""Return the AMD pip index URL for the given GPU arch, or None if unsupported."""
arch_family = _GFX_TO_AMD_INDEX_ARCH.get(gfx_arch or "")
@ -1647,22 +1783,24 @@ def _ensure_rocm_torch() -> None:
# An explicit ROCm pin commits to ROCm wheels regardless of the visible GPU (headless / CI).
# Mirror _ensure_cuda_torch: skip the NVIDIA/no-AMD/unreadable gates.
_rocm_pin = _explicit_rocm_torch_index_url()
_inferred_linux_gfx = (
_infer_linux_amd_gfx_arch() if (_rocm_pin is None and not IS_WINDOWS) else None
)
if _rocm_pin is None:
# NVIDIA takes precedence on mixed hosts (only if a GPU is usable).
if _has_usable_nvidia_gpu():
return
# _has_rocm_gpu() (rocminfo / amd-smi rows) is the authoritative AMD-host signal;
# the old /opt/rocm-or-hipcc gate broke runtime-only ROCm installs.
if not _has_rocm_gpu():
if not _has_rocm_gpu() and not _inferred_linux_gfx:
return # no AMD GPU visible
ver = _detect_rocm_version()
if ver is None:
if _rocm_pin is None:
if _rocm_pin is None and not _inferred_linux_gfx:
print(" ROCm detected but version unreadable -- skipping torch reinstall")
return
# Explicit pin: the pinned leaf drives the install, so an unreadable host version
# is fine (sentinel keeps ver comparisons defined).
# Explicit pin or inferred gfx: the index drives the install.
ver = (0, 0)
# Probe whether torch links against HIP, capturing the installed ROCm tag for pin-mismatch
@ -1712,6 +1850,44 @@ def _ensure_rocm_torch() -> None:
rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch
# Inferred-gfx path: ROCm runtime missing but install.sh would route to AMD wheels.
# Gated on the runtime NOT enumerating a GPU: when it can, the runtime-visible
# arch (Strix override / generic below) decides, not cpuinfo -- a mixed Strix
# APU + dGPU box with HIP_VISIBLE_DEVICES on the dGPU must not get APU wheels.
# An explicit UNSLOTH_ROCM_GFX_ARCH is exempt from that runtime gate (mirrors
# install.sh): a visible GPU with an unreadable/unsupported ROCm version must
# not silently discard the user's named arch and leave CPU torch in place.
_gfx_override_env = (os.environ.get("UNSLOTH_ROCM_GFX_ARCH") or "").strip().lower()
if (
_inferred_linux_gfx
and not has_hip_torch
and _rocm_pin is None
and (_gfx_override_env or not _has_rocm_gpu())
):
index_url = _amd_arch_index_url(_inferred_linux_gfx)
if index_url is not None:
_torch_pkg, _vision_pkg, _audio_pkg = _WINDOWS_ROCM_TORCH_PKG_SPECS.get(
_inferred_linux_gfx, ("torch", "torchvision", "torchaudio")
)
print(
f"\n {_inferred_linux_gfx} inferred (ROCm runtime not visible) -- "
f"installing torch from {_strip_index_url_credentials(index_url)}\n"
f" AMD wheels bundle their own ROCm runtime; install the kernel stack "
f"for native GPU compute.\n"
)
pip_install(
f"ROCm torch (inferred {_inferred_linux_gfx})",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
)
rocm_torch_ready = True
# 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.
@ -1776,8 +1952,11 @@ def _ensure_rocm_torch() -> None:
constrain = False,
)
rocm_torch_ready = True
elif not has_hip_torch or _rocm_pin_mismatch:
elif not rocm_torch_ready:
# Reinstall when torch is not ROCm yet, OR a ROCm build's family differs from a pin.
# Gate on rocm_torch_ready (not has_hip_torch alone) so a successful inferred-gfx
# install above is not overwritten by the generic pytorch.org/rocmX.Y path -- that
# would undo the fresh-ROCm/no-/dev/kfd repair this path exists for (Codex P1 #7305).
# Honour a ROCm pin verbatim; else pick the newest wheel tag <= host.
_override_idx = _explicit_rocm_torch_index_url()
if _override_idx is not None:

View file

@ -9,6 +9,7 @@ import subprocess
import sys
import tempfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, mock_open, patch, PropertyMock
import pytest
@ -560,9 +561,13 @@ class TestDetectRocmVersion:
class TestEnsureRocmTorch:
"""Verify ROCm torch reinstall logic."""
# _infer_linux_amd_gfx_arch mocked to None: on a real Strix host the live
# /proc/cpuinfo would otherwise take the inferred-install path and break
# these "must not install" hosts (environment leak, not the code under test).
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
def test_no_rocm_skips(self, mock_nvidia, mock_pip):
@patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = None)
def test_no_rocm_skips(self, mock_infer, mock_nvidia, mock_pip):
"""No ROCm toolchain should skip entirely."""
# Pin _detect_windows_gfx_arch to None so a real AMD test host's WMI
# fallback can't defeat the "no ROCm anywhere" premise.
@ -572,6 +577,105 @@ class TestEnsureRocmTorch:
_ensure_rocm_torch()
mock_pip.assert_not_called()
@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 = False)
@patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = "gfx1151")
@patch.object(stack_mod, "_detect_rocm_version", return_value = None)
def test_inferred_gfx_without_rocm_runtime_installs_amd_index(
self, mock_ver, mock_infer, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""Strix Halo without /dev/kfd must still get AMD gfx1151 wheels (unslothai#7301)."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"|2.10.0+cpu\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 "gfx1151" 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")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = False)
@patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = "gfx1151")
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = [])
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_inferred_gfx_not_overwritten_when_rocm_userland_readable(
self, mock_ver, mock_gfx, mock_infer, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""Codex P1 #7305: after an inferred per-arch install, do not fall through to the
generic pytorch.org/rocmX.Y reinstall just because has_hip_torch is still False.
Readable ROCm userland without /dev/kfd is exactly the case that used to overwrite
the AMD gfx wheels."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"|2.10.0+cpu\n"
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
assert mock_pip.call_count == 1, mock_pip.call_args_list
torch_call = str(mock_pip.call_args_list[0])
assert "gfx1151" in torch_call
assert "rocm7.1" not in torch_call
assert "download.pytorch.org" not 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")
@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, "_infer_linux_amd_gfx_arch", return_value = "gfx1151")
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1100"])
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_inference_yields_to_runtime_visible_gpu(
self, mock_ver, mock_gfx, mock_infer, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""When the runtime CAN enumerate a GPU, the cpuinfo inference must not
install wheels: a mixed Strix APU + dGPU box with the dGPU selected would
otherwise get gfx1151 wheels for a gfx1100 GPU. The runtime-visible arch
(Strix override / generic branch) decides instead."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"|2.10.0+cpu\n"
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
all_calls = str(mock_pip.call_args_list) + str(mock_pip_try.call_args_list)
assert "gfx1151" not in all_calls, all_calls
assert "rocm7.1" in all_calls, all_calls
@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_amd_gfx_codes", return_value = [])
@patch.object(stack_mod, "_detect_rocm_version", return_value = None)
def test_gfx_override_installs_despite_visible_rocm(
self, mock_ver, mock_gfx, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""#7305 review: an explicit UNSLOTH_ROCM_GFX_ARCH is exempt from the
not-_has_rocm_gpu() gate (mirrors install.sh). A visible GPU with an
unreadable ROCm version must not silently discard the user's named arch
and leave CPU torch in place -- the per-arch install runs."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"|2.10.0+cpu\n"
with patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": "gfx1151"}):
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
assert mock_pip.call_count == 1, mock_pip.call_args_list
torch_call = str(mock_pip.call_args_list[0])
assert "gfx1151" in torch_call
assert "download.pytorch.org" not 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")
@ -683,9 +787,10 @@ class TestEnsureRocmTorch:
@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, "_infer_linux_amd_gfx_arch", return_value = None)
@patch.object(stack_mod, "_detect_rocm_version", return_value = None)
def test_version_unreadable_prints_warning(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, capsys
self, mock_ver, mock_infer, mock_gpu, mock_nvidia, mock_pip, capsys
):
"""ROCm detected but version unreadable should print warning and skip."""
with patch("os.path.isdir", return_value = True):
@ -1042,7 +1147,8 @@ class TestEnsureRocmTorch:
@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 = False)
def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip):
@patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = None)
def test_no_gpu_with_rocm_tools_skips(self, mock_infer, mock_gpu, mock_nvidia, mock_pip):
"""ROCm tools present but no actual AMD GPU should skip entirely."""
# Pin the Windows arch probe to None so a real AMD host's WMI fallback
# can't defeat the "no actual GPU" premise.
@ -2122,6 +2228,7 @@ class TestGfxArchNameFallback:
"name, expected",
[
("AMD Radeon(TM) 8060S Graphics", "gfx1151"),
("AMD Radeon(TM) 8065S Graphics", "gfx1151"),
("AMD Ryzen AI MAX+ 395 w/ Radeon 8060S", "gfx1151"),
("AMD Radeon(TM) 890M", "gfx1150"),
("AMD Ryzen AI 9 HX 370 w/ Radeon 890M", "gfx1150"),
@ -3189,6 +3296,286 @@ _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh"
class TestStrixRocm71Override:
"""install.sh routes gfx1151/gfx1150 to AMD's arch index instead of ROCm 7.1 (_grouped_mm segfault)."""
def test_linux_gfx_inference_helpers_present(self):
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
assert "_infer_linux_amd_gfx_arch" in source
assert "_amd_arch_index_family_for_gfx" in source
assert "_amd_gpu_present_via_pci" in source
assert "unslothai#7301" in source
def test_infer_linux_amd_gfx_from_cpuinfo(self):
assert stack_mod._linux_amd_gfx_from_cpuinfo is not None
with patch.object(
Path,
"read_text",
return_value = "model name : AMD Ryzen AI Max+ 395 w/ Radeon 8060S\n",
):
assert stack_mod._linux_amd_gfx_from_cpuinfo() == "gfx1151"
# 8065S (Gorgon Halo) must match on the Radeon name alone, even without the
# "Ryzen AI Max" branding (mirrors setup.sh / setup.ps1 which list 8065S).
with patch.object(Path, "read_text", return_value = "model name : AMD Radeon 8065S\n"):
assert stack_mod._linux_amd_gfx_from_cpuinfo() == "gfx1151"
def test_infer_gfx_gated_out_of_wsl_without_runtime(self):
"""On WSL the cpuinfo/lspci inference must be skipped unless the WSL ROCDXG
runtime (librocdxg) is present: a bare `unsloth studio update` must not
install per-arch ROCm wheels into an env that still can't expose the GPU.
An explicit UNSLOTH_ROCM_GFX_ARCH override stays authoritative regardless."""
m = stack_mod
with (
patch.object(m, "_linux_amd_gfx_from_cpuinfo", return_value = "gfx1151"),
patch.object(m, "_linux_amd_gfx_from_lspci", return_value = None),
# PCI evidence present (the WSL branch never consults it anyway).
patch.object(m, "_linux_amd_display_device_present", return_value = True),
patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": ""}),
):
# WSL + no runtime -> inference suppressed (CPU torch stays).
with (
patch.object(m, "_is_wsl", return_value = True),
patch.object(m, "_wsl_rocm_runtime_present", return_value = False),
):
assert m._infer_linux_amd_gfx_arch() is None
# WSL + runtime present (this dev box) -> inference still runs.
with (
patch.object(m, "_is_wsl", return_value = True),
patch.object(m, "_wsl_rocm_runtime_present", return_value = True),
):
assert m._infer_linux_amd_gfx_arch() == "gfx1151"
# Native Linux (not WSL) -> the gate never applies.
with (
patch.object(m, "_is_wsl", return_value = False),
patch.object(m, "_wsl_rocm_runtime_present", return_value = False),
):
assert m._infer_linux_amd_gfx_arch() == "gfx1151"
# Explicit override wins even on a bare WSL box (no runtime).
with (
patch.object(m, "_is_wsl", return_value = True),
patch.object(m, "_wsl_rocm_runtime_present", return_value = False),
patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": "gfx1151"}),
):
assert m._infer_linux_amd_gfx_arch() == "gfx1151"
def test_infer_gfx_requires_amd_display_device_on_native_linux(self):
"""A VM/container on a Strix host still shows the host CPU model in
/proc/cpuinfo while receiving no AMD GPU, so on native Linux the
CPU-model inference must require an AMD PCI display device (#7305
review). WSL is exempt (no PCI enumeration there; the librocdxg gate is
the evidence) and the explicit override stays authoritative."""
m = stack_mod
with (
patch.object(m, "_linux_amd_gfx_from_cpuinfo", return_value = "gfx1151"),
patch.object(m, "_linux_amd_gfx_from_lspci", return_value = None),
patch.object(m, "_is_wsl", return_value = False),
patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": ""}),
):
# No AMD display device -> the CPU-model text alone must not infer.
with patch.object(m, "_linux_amd_display_device_present", return_value = False):
assert m._infer_linux_amd_gfx_arch() is None
# Device present -> inference unchanged.
with patch.object(m, "_linux_amd_display_device_present", return_value = True):
assert m._infer_linux_amd_gfx_arch() == "gfx1151"
# Explicit override needs no device evidence (headless/cross-install).
with (
patch.object(m, "_is_wsl", return_value = False),
patch.object(m, "_linux_amd_display_device_present", return_value = False),
patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": "GFX1151"}),
):
assert m._infer_linux_amd_gfx_arch() == "gfx1151"
def test_install_sh_cpuinfo_inference_requires_pci_evidence(self):
"""install.sh mirror of the VM/container guard: both cpuinfo greps must be
gated on _gpu_evidence (AMD PCI display device via _amd_gpu_present_via_pci,
or the WSL librocdxg gate), and the gate must sit before the first grep."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
body = _extract_sh_function_body(source, "_infer_linux_amd_gfx_arch")
assert body, "could not extract _infer_linux_amd_gfx_arch"
pci = body.find("_amd_gpu_present_via_pci")
infer = body.find("grep -qiE 'Ryzen AI Max")
assert pci >= 0 and infer >= 0
assert pci < infer, "the PCI evidence check must run before the cpuinfo inference"
assert (
body.count('[ -n "$_gpu_evidence" ] && grep -qiE') == 2
), "both cpuinfo greps (gfx1151 and gfx1150) must be gated on _gpu_evidence"
def test_lspci_scan_covers_all_display_controllers(self):
"""The lspci fallback must scan every display-class line, not just the
first: a non-AMD controller (Intel iGPU, ASPEED BMC) often enumerates
before the AMD dGPU. Non-AMD vendors must never map (an NVIDIA GeForce
GTX 860M would otherwise hit the AMD 860M pattern), and a 0000: PCI
domain prefix must not break matching."""
m = stack_mod
def fake_lspci(stdout):
result = SimpleNamespace(returncode = 0, stdout = stdout)
return (
patch.object(m.shutil, "which", return_value = "/usr/bin/lspci"),
patch.object(m.subprocess, "run", return_value = result),
)
intel_then_amd = (
"00:02.0 VGA compatible controller [0300]: Intel Corporation Raptor Lake-S GT1 [8086:a780]\n"
"03:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. [AMD/ATI]"
" Navi 31 [Radeon RX 7900 XT] [1002:744c]\n"
)
nvidia_only = "01:00.0 3D controller [0302]: NVIDIA Corporation GM107M [GeForce GTX 860M] [10de:1392]\n"
domain_prefixed = (
"0000:c5:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. [AMD/ATI]"
" Strix Halo [Radeon Graphics / Radeon 8060S] [1002:150e]\n"
)
unmapped_then_mapped = (
"03:00.0 Display controller [0380]: Advanced Micro Devices, Inc. [AMD/ATI]"
" Cape Verde [FirePro W600] [1002:6821]\n"
"04:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. [AMD/ATI]"
" Navi 33 [Radeon RX 7600] [1002:7480]\n"
)
for stdout, expected in (
(intel_then_amd, "gfx1100"),
(nvidia_only, None),
(domain_prefixed, "gfx1151"),
(unmapped_then_mapped, "gfx1102"),
):
w, r = fake_lspci(stdout)
with w, r:
assert m._linux_amd_gfx_from_lspci() == expected, stdout
def test_install_sh_lspci_scan_covers_all_display_controllers(self):
"""install.sh mirror of the scan-all behaviour, executed with a shimmed
lspci: Intel-first still finds the AMD dGPU, NVIDIA-only maps nothing
(860M collision), a domain-prefixed AMD line still maps."""
shell = shutil.which("bash")
if not shell:
pytest.skip("bash needed to execute the probe block")
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
name_fn = re.search(
r"^_infer_amd_gfx_arch_from_gpu_name\(\) \{\n.*?\n\}\n", source, re.S | re.M
)
scan = re.search(
r"^ if command -v lspci[^\n]*\n.*?\nEOF\n fi\n return 1\n", source, re.S | re.M
)
assert name_fn and scan, "could not extract the lspci scan block"
cases = (
(
"00:02.0 VGA compatible controller [0300]: Intel Corporation UHD [8086:a780]\n"
"03:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. [AMD/ATI]"
" Navi 31 [Radeon RX 7900 XT] [1002:744c]",
"OK:gfx1100",
),
(
"01:00.0 3D controller [0302]: NVIDIA Corporation GM107M [GeForce GTX 860M] [10de:1392]",
"OK:",
),
(
"0000:c5:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc."
" [AMD/ATI] Strix Halo [Radeon 8060S] [1002:150e]",
"OK:gfx1151",
),
)
for lspci_out, expected in cases:
with tempfile.TemporaryDirectory() as d:
p = os.path.join(d, "lspci")
with open(p, "w", encoding = "utf-8") as f:
f.write(f'#!/bin/sh\ncat <<"EOT"\n{lspci_out}\nEOT\n')
os.chmod(p, 0o755)
script = (
"set -euo pipefail\n"
+ name_fn.group(0)
+ "probe() {\n"
+ scan.group(0)
+ "}\nprintf 'OK:%s\\n' \"$(probe || true)\"\n"
)
env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""))
r = subprocess.run([shell, "-c", script], env = env, capture_output = True, text = True)
assert r.returncode == 0, f"scan aborted: {r.stderr}"
assert (
r.stdout.splitlines()[-1] == expected
), f"lspci scan wrong for {lspci_out!r}: {r.stdout!r}"
def test_install_sh_infer_gfx_gated_on_wsl_runtime(self):
"""install.sh's _infer_linux_amd_gfx_arch must, like the Python side, skip
the cpuinfo/lspci inference on WSL unless librocdxg is present -- the
override still returns first, so it stays authoritative."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
body = _extract_sh_function_body(source, "_infer_linux_amd_gfx_arch")
assert body, "could not extract _infer_linux_amd_gfx_arch"
override = body.find("UNSLOTH_ROCM_GFX_ARCH")
dxg = body.find("/dev/dxg")
rocdxg = body.find("librocdxg")
# Anchor on the first cpuinfo *inference* (the grep), not a comment mention.
infer = body.find("grep -qiE 'Ryzen AI Max")
assert override >= 0 and dxg >= 0 and rocdxg >= 0 and infer >= 0
assert "microsoft" in body, "WSL gate must also detect WSL via /proc/version"
assert override < dxg, "the explicit override must return before the WSL gate"
assert (
dxg < infer and rocdxg < infer
), "the WSL/librocdxg gate must run before the cpuinfo/lspci inference"
def test_install_sh_reroute_is_x86_64_only(self):
"""The Linux inferred-gfx reroute must be x86_64-only: ROCm torch wheels are
not published for arm64, so an inferred/overridden gfx must not push an
arm64 host to the AMD arch index (get_torch_index_url returns CPU there)."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
idx = source.find("_linux_inferred_gfx=$(_infer_linux_amd_gfx_arch")
assert idx >= 0, "reroute consumer not found"
window = source[max(0, idx - 400) : idx]
assert (
'case "$_ARCH" in x86_64|amd64)' in window
), "the inferred-gfx reroute must guard on x86_64|amd64 arch"
def test_install_sh_reroute_skips_visible_rocm_gpu(self):
"""A */cpu index on a host whose AMD GPU IS visible to the ROCm probes is a
deliberate fallback (unsupported/unreadable ROCm version, warned about in
get_torch_index_url), not a missing runtime: the reroute must not override
it with inferred per-arch wheels. The explicit UNSLOTH_ROCM_GFX_ARCH
override must still win either way."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
idx = source.find("_linux_inferred_gfx=$(_infer_linux_amd_gfx_arch")
assert idx >= 0, "reroute consumer not found"
window = source[max(0, idx - 700) : idx]
assert (
"! _has_amd_rocm_gpu" in window
), "the reroute must be gated on _has_amd_rocm_gpu being false"
assert (
'[ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu' in window
), "an explicit UNSLOTH_ROCM_GFX_ARCH override must bypass the visible-GPU gate"
def test_install_sh_reroute_exports_gfx_for_setup_sh(self):
"""The inferred arch must be exported as UNSLOTH_ROCM_GFX_ARCH so the
downstream setup.sh run (which re-probes ROCm independently and finds
nothing on these runtime-less hosts) routes llama.cpp to the matching
ROCm prebuilt instead of the CPU one -- setup.sh and
install_llama_prebuilt.py both read that env var."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
assign = source.find('TORCH_INDEX_URL="${_amd_mirror}/${_amd_family}/"')
assert assign >= 0, "inferred-gfx index assignment not found"
block_end = source.find("esac", assign)
assert (
'export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx"' in source[assign:block_end]
), "the reroute must export the inferred gfx for the setup.sh handoff"
# setup.sh's side of the handoff must still exist.
setup_source = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8")
assert "UNSLOTH_ROCM_GFX_ARCH" in setup_source
def test_amd_arch_index_url_linux_honors_amd_mirror(self):
"""On Linux the inferred-gfx repair must honour UNSLOTH_AMD_ROCM_MIRROR (the
var install.sh uses), not the Windows mirror var, so a mirrored/air-gapped
Linux install does not silently fall back to repo.amd.com. Windows still
delegates to the Windows mirror path."""
m = stack_mod
with (
patch.object(m, "IS_WINDOWS", False),
patch.dict(os.environ, {"UNSLOTH_AMD_ROCM_MIRROR": "https://mirror.local/rocm"}),
):
assert m._amd_arch_index_url("gfx1151") == "https://mirror.local/rocm/gfx1151/"
with (
patch.object(m, "IS_WINDOWS", False),
patch.dict(os.environ, {"UNSLOTH_AMD_ROCM_MIRROR": ""}),
):
assert m._amd_arch_index_url("gfx1151") == "https://repo.amd.com/rocm/whl/gfx1151/"
assert m._amd_arch_index_url("gfx9999") is None
# Windows path is unchanged: delegate to the Windows mirror helper.
with patch.object(m, "IS_WINDOWS", True):
assert m._amd_arch_index_url("gfx1151") == m._windows_rocm_index_url("gfx1151")
def test_strix_gfx_detection_in_install_sh(self):
"""install.sh must detect gfx1151 and gfx1150 for the override."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")