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:
parent
c267895538
commit
978ae4745b
3 changed files with 714 additions and 8 deletions
140
install.sh
140
install.sh
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue