Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 84 additions & 0 deletions tests/_zoo_rocm_spoof.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""ROCm/RDNA spoof: present torch as an AMD Radeon (RDNA 2/3/4) card on a
GPU-less host, so hip paths (device_type -> "hip", llama.cpp ROCm bundle) are
testable in CPU-only CI with no AMD hardware. The ROCm sibling of
_zoo_aggressive_cuda_spoof.py: it reuses that spoof's torch.cuda no-op machinery
and overlays the AMD identity (torch.version.hip, gcnArchName, Radeon name).
Apply BEFORE importing unsloth/unsloth_zoo, since DEVICE_TYPE is cached there.
"""

from __future__ import annotations

import importlib.util
import os
import sys

# gfx -> (marketing name, (capability major, minor), torch.version.hip). hip is
# the ROCm build torch was made against (RDNA2/3 ship 6.x; gfx1102/115x/RDNA4 7.2).
_PROFILES: dict[str, tuple[str, tuple[int, int], str]] = {
"gfx1030": ("AMD Radeon RX 6900 XT", (10, 3), "6.4.43483"), # RDNA2
"gfx1031": ("AMD Radeon RX 6700 XT", (10, 3), "6.4.43483"),
"gfx1032": ("AMD Radeon RX 6600", (10, 3), "6.4.43483"),
"gfx1034": ("AMD Radeon RX 6400", (10, 3), "6.4.43483"),
"gfx1100": ("AMD Radeon RX 7900 XTX", (11, 0), "6.4.43483"), # RDNA3
"gfx1101": ("AMD Radeon RX 7800 XT", (11, 0), "6.4.43483"),
"gfx1102": ("AMD Radeon RX 7600", (11, 0), "7.2.1"),
"gfx1150": ("AMD Radeon 890M", (11, 5), "7.2.1"), # RDNA3.5 APU
"gfx1151": ("AMD Radeon 8060S", (11, 5), "7.2.1"),
"gfx1200": ("AMD Radeon RX 9060 XT", (12, 0), "7.2.1"), # RDNA4
"gfx1201": ("AMD Radeon RX 9070 XT", (12, 0), "7.2.1"),
}


def _cuda_spoof():
"""Load the sibling CUDA spoof by path (robust to sys.path), so we reuse its
torch.cuda machinery instead of duplicating it."""
if "_zoo_aggressive_cuda_spoof" in sys.modules:
return sys.modules["_zoo_aggressive_cuda_spoof"]
path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_zoo_aggressive_cuda_spoof.py")
spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
Comment on lines +40 to +42

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Defensively check if spec or spec.loader is None before attempting to load the module. If either is None, raising a clear ImportError is much more helpful for debugging than letting it fail with a generic AttributeError or ValueError.

Suggested change
spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path)
if spec is None or spec.loader is None:
raise ImportError(f"Could not load spec for _zoo_aggressive_cuda_spoof at {path}")
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)

sys.modules["_zoo_aggressive_cuda_spoof"] = mod
return mod


def apply(gfx: str = "gfx1100", device_count: int = 1) -> None:
"""Present torch as `gfx`. Re-callable to switch arch (identity is overlaid;
the underlying no-op machinery is applied once)."""
import torch

if gfx not in _PROFILES:
raise KeyError(f"Unknown gfx {gfx!r}; known: {', '.join(_PROFILES)}")
name, cap, hip = _PROFILES[gfx]

_cuda_spoof().apply() # is_available/device_count/streams/rng/amp/...

# Overlay the AMD identity on top of the (NVIDIA-shaped) CUDA spoof.
torch.version.hip = hip
torch.version.cuda = None
torch.cuda.device_count = lambda: device_count
torch.cuda.get_device_name = lambda *a, **k: name
torch.cuda.get_device_capability = lambda *a, **k: cap
torch.cuda.get_arch_list = lambda: [gfx]

class _Props:
pass

_p = _Props()
_p.name = name
_p.gcnArchName = f"{gfx}:sramecc-:xnack-" # ROCm advertises feature flags
_p.major, _p.minor = cap
_p.total_memory = 16 * 1024**3
_p.multi_processor_count = 40
_p.warp_size = 32 # RDNA wavefront (CDNA is 64)
_p.is_integrated = gfx in ("gfx1150", "gfx1151")
_p.is_multi_gpu_board = False
torch.cuda.get_device_properties = lambda *a, **k: _p


if __name__ == "__main__":
apply()
import torch
print("ROCm spoof applied:", torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName)
84 changes: 84 additions & 0 deletions tests/studio/install/test_rocm_rdna_routing.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""RDNA 2/3/4 routing, validated on CPU-only CI with no AMD hardware.

tests/_zoo_rocm_spoof.py presents torch as each Radeon gfx arch, then we assert
unsloth_zoo routes it: device_type -> "hip", llama.cpp target -> ("rocm", gfx),
and the per-family ROCm bundle suffix. The torch-facing checks run in a
subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached
at import) resolves from a clean process.
"""

from __future__ import annotations

import json
import subprocess
import sys
from pathlib import Path

import pytest

pytest.importorskip("torch")
pytest.importorskip("unsloth_zoo")

_TESTS_DIR = Path(__file__).resolve().parents[2] # tests/

# gfx -> (expected llama.cpp target, expected ROCm bundle family).
_ARCHES = {
"gfx1030": (("rocm", "gfx1030"), "gfx103X"), # RDNA2
"gfx1031": (("rocm", "gfx1031"), "gfx103X"),
"gfx1032": (("rocm", "gfx1032"), "gfx103X"),
"gfx1034": (("rocm", "gfx1034"), "gfx103X"),
"gfx1100": (("rocm", "gfx1100"), "gfx110X"), # RDNA3
"gfx1101": (("rocm", "gfx1101"), "gfx110X"),
"gfx1102": (("rocm", "gfx1102"), "gfx110X"),
"gfx1150": (("rocm", "gfx1150"), "gfx1150"), # RDNA3.5 APU (self-family)
"gfx1151": (("rocm", "gfx1151"), "gfx1151"),
"gfx1200": (("rocm", "gfx1200"), "gfx120X"), # RDNA4
"gfx1201": (("rocm", "gfx1201"), "gfx120X"),
}

# Child: spoof each arch, then record device_type once (fresh import) and the
# live llama.cpp target per arch. Emits one JSON line the parent parses.
_CHILD = """
import json, sys
sys.path.insert(0, {tests!r})
import _zoo_rocm_spoof as spoof
arches = {arches!r}
spoof.apply(arches[0])
from unsloth_zoo.device_type import get_device_type, is_hip
device_type = [get_device_type(), is_hip()]
from unsloth_zoo import llama_cpp as lc
targets = {{}}
for gfx in arches:
spoof.apply(gfx)
targets[gfx] = list(lc._detect_gpu_target())
print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}}))
"""
Comment on lines +43 to +57

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To make the test runner robust to where pytest is invoked from (especially if run from a subdirectory or if unsloth_zoo is not globally installed in the environment), explicitly insert the repository root directory (_TESTS_DIR.parent) into the child process's sys.path.

Suggested change
_CHILD = """
import json, sys
sys.path.insert(0, {tests!r})
import _zoo_rocm_spoof as spoof
arches = {arches!r}
spoof.apply(arches[0])
from unsloth_zoo.device_type import get_device_type, is_hip
device_type = [get_device_type(), is_hip()]
from unsloth_zoo import llama_cpp as lc
targets = {{}}
for gfx in arches:
spoof.apply(gfx)
targets[gfx] = list(lc._detect_gpu_target())
print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}}))
"""
_CHILD = """
import json, sys
sys.path.insert(0, {tests!r})
sys.path.insert(0, {repo!r})
import _zoo_rocm_spoof as spoof
arches = {arches!r}
spoof.apply(arches[0])
from unsloth_zoo.device_type import get_device_type, is_hip
device_type = [get_device_type(), is_hip()]
from unsloth_zoo import llama_cpp as lc
targets = {{}}
for gfx in arches:
spoof.apply(gfx)
targets[gfx] = list(lc._detect_gpu_target())
print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}}))
"""



@pytest.fixture(scope = "module")
def routed():
code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES))
proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True)
line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None)
assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
return json.loads(line[len("RESULT ") :])


@pytest.mark.parametrize("gfx", list(_ARCHES))
def test_detect_gpu_target(routed, gfx):
# RDNA card is routed to its ROCm gfx target (drives the llama.cpp bundle).
assert tuple(routed["targets"][gfx]) == _ARCHES[gfx][0]


def test_device_type_is_hip(routed):
# An RDNA card must resolve the compute device_type to "hip".
assert routed["device_type"] == ["hip", True]


@pytest.mark.parametrize("gfx", list(_ARCHES))
def test_rocm_gfx_family(gfx):
# Pure mapping (no torch): each gfx picks the right per-family ROCm bundle.
from unsloth_zoo import llama_cpp as lc
assert lc._rocm_gfx_family(gfx) == _ARCHES[gfx][1]
Loading