# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Unit tests for _rocm_classify_unified_memory (ROCm OOM-guard classifier). Three paths: (1) canonical gcnArchName, (2) alternate-spelling attr, (3) all arch attrs absent -> device-name substring match. Regression: Strix Halo (gfx1151) was misclassified as discrete on Radeon wheels that set props.name="Radeon 8060S Graphics" but no gcnArchName, applying the wrong headroom factor on a 128 GiB unified-memory pool. """ from __future__ import annotations from pathlib import Path from types import SimpleNamespace import pytest from core.training.worker import _rocm_classify_unified_memory # ── helpers ────────────────────────────────────────────────────────────────── def _props(**kwargs) -> SimpleNamespace: """Build a fake device-properties object with the given attributes.""" return SimpleNamespace(**kwargs) # ── Path 0: props.is_integrated (driver's own unified-memory answer) ───────── class TestIsIntegratedSignal: """hipDeviceProp_t.integrated wins when truthy; 0/absent never downgrades. Same universal gate PR #5988's UMA safetensors fast-load uses -- keeps Studio's two unified-memory consumers on one signal.""" def test_integrated_upgrades_unknown_apu(self) -> None: # gfx1103 Phoenix iGPU: outside the hardcoded arch set, but the # driver says integrated -> unified. props = _props(gcnArchName = "gfx1103", name = "Radeon 780M", is_integrated = 1) gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "gfx1103" assert is_unified is True def test_integrated_wins_without_any_arch(self) -> None: props = _props(name = "Some Future APU", is_integrated = 1) gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "" assert is_unified is True def test_zero_does_not_downgrade_known_apu(self) -> None: # A wheel that zeroes the field must not flip Strix Halo to discrete. props = _props(gcnArchName = "gfx1151", name = "x", is_integrated = 0) gcn, is_unified = _rocm_classify_unified_memory(props) assert is_unified is True def test_absent_keeps_existing_behavior(self) -> None: props = _props(gcnArchName = "gfx1201", name = "RX 9070 XT") gcn, is_unified = _rocm_classify_unified_memory(props) assert is_unified is False def test_discrete_with_zero_stays_discrete(self) -> None: props = _props(gcnArchName = "gfx1100", name = "RX 7900 XTX", is_integrated = 0) gcn, is_unified = _rocm_classify_unified_memory(props) assert is_unified is False # ── Path 1: canonical gcnArchName ──────────────────────────────────────────── class TestCanonicalGcnArchName: """gcnArchName is present and populated.""" @pytest.mark.parametrize( "arch, expected_unified", [ ("gfx1150", True), # Strix Point ("gfx1151", True), # Strix Halo ("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete ("gfx906", False), # MI50 — discrete server GPU ("gfx1201", False), # RX 9070 XT — discrete ], ) def test_canonical_attr(self, arch: str, expected_unified: bool) -> None: props = _props(gcnArchName = arch, name = "irrelevant") gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == arch assert is_unified is expected_unified def test_arch_with_colon_suffix_stripped(self) -> None: """gcnArchName can carry xnack/sramecc suffix; only the base is kept.""" props = _props(gcnArchName = "gfx1151:xnack-", name = "irrelevant") gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "gfx1151" assert is_unified is True def test_canonical_attr_wins_over_name(self) -> None: """Arch attr takes priority; device name is ignored.""" # Discrete arch, but name looks like a unified SKU — arch must win. props = _props(gcnArchName = "gfx1100", name = "Radeon 890M") gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "gfx1100" assert is_unified is False # ── Path 2: alternate-spelling fallback ────────────────────────────────────── class TestAlternateSpellingFallback: """gcnArchName is missing but an alternate attr spelling is present.""" @pytest.mark.parametrize( "attr_name", ["gcn_arch_name", "arch_name", "gfx_arch_name"], ) def test_alternate_attr_unified(self, attr_name: str) -> None: props = _props(**{attr_name: "gfx1151"}, name = "Radeon 8060S Graphics") gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "gfx1151" assert is_unified is True @pytest.mark.parametrize( "attr_name", ["gcn_arch_name", "arch_name", "gfx_arch_name"], ) def test_alternate_attr_discrete(self, attr_name: str) -> None: props = _props(**{attr_name: "gfx1201"}, name = "Radeon RX 9070 XT") gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "gfx1201" assert is_unified is False def test_first_non_empty_attr_wins(self) -> None: """With multiple alternate attrs, the first non-empty one wins.""" props = _props(gcn_arch_name = "gfx1151", arch_name = "gfx1100", name = "irrelevant") gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "gfx1151" assert is_unified is True # ── Path 3: device-name fallback ───────────────────────────────────────────── class TestDeviceNameFallback: """ALL arch attrs absent — classifier must rely solely on device name.""" # --- unified-memory devices that MUST be detected --- @pytest.mark.parametrize( "device_name", [ # gfx1150 Strix Point "Radeon 890M", "AMD Radeon 890M Graphics", "RADEON 890M", # case-insensitive "Radeon 880M", "AMD Radeon 880M Graphics", # gfx1151 Strix Halo — the regression case from the review "Radeon 8060S Graphics", # Ryzen AI MAX+ 395 (as returned by torch) "AMD Radeon 8060S", "Radeon 8050S Graphics", # cut-down Strix Halo SKU "AMD Radeon 8050S", # case variants "RADEON 8060S GRAPHICS", "radeon 8050s", ], ) def test_unified_memory_detected(self, device_name: str) -> None: props = _props(name = device_name) gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "", f"expected empty gcn_arch, got {gcn!r}" assert is_unified is True, f"device {device_name!r} should be classified as unified-memory" # --- discrete devices that must NOT be mis-classified --- @pytest.mark.parametrize( "device_name", [ "Radeon RX 9070 XT", "AMD Radeon RX 7900 XTX", "Radeon RX 6900 XT", "Radeon Pro W7900", "AMD Instinct MI300X", # Superficially similar substrings but discrete "Radeon RX 580", "Radeon VII", ], ) def test_discrete_not_misclassified(self, device_name: str) -> None: props = _props(name = device_name) gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "" assert ( is_unified is False ), f"discrete device {device_name!r} should NOT be classified as unified-memory" def test_empty_name_returns_false(self) -> None: """Absent name must not crash and must default to discrete.""" props = _props() # no 'name' attr at all gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "" assert is_unified is False def test_none_name_returns_false(self) -> None: props = _props(name = None) gcn, is_unified = _rocm_classify_unified_memory(props) assert gcn == "" assert is_unified is False # ── Fraction selection (source-pinned) ─────────────────────────────────────── _WORKER_PY = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" class TestMemFractionSelection: """Pin the per-platform fraction policy in worker.py section 1g. On native Windows, torch.cuda.mem_get_info's total is the WDDM budget the driver grants HIP -- the OS share of RAM is already outside it, so a 0.80 cap double-taxes (field report: 48.49 GiB budget -> '38.79 GiB allowed' OOM denying a 47.29 GiB load that fit in free memory). 1.0 removes the double-tax; current AMD Windows wheels enforce only sub-1.0 fractions, so it behaves like torch's uncapped default with WDDM arbitrating residency (measured on gfx1151).""" def test_unified_win32_uses_budget_exact_fraction(self) -> None: source = _WORKER_PY.read_text(encoding = "utf-8") assert '1.0 if sys.platform == "win32" else 0.80' in source def test_discrete_keeps_090(self) -> None: source = _WORKER_PY.read_text(encoding = "utf-8") assert "_mem_fraction = 0.90" in source def test_win32_unified_logs_vgm_hint(self) -> None: """Users must learn the WDDM budget is raisable (BIOS UMA / AMD Software Variable Graphics Memory) instead of assuming a bug.""" source = _WORKER_PY.read_text(encoding = "utf-8") assert "Variable Graphics Memory" in source