From b4921e2a835b08f2ec5fe26fec39897bc5f18041 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 9 Jun 2026 21:58:28 -0700 Subject: [PATCH] feat(studio): Spark unified-memory OOM guard in the training worker (Strix Halo parity) PR #5301 protects ROCm unified-memory APUs (Strix Halo gfx1150/gfx1151) with a default set_per_process_memory_fraction(0.80) at training-worker startup, because exhausting a shared GPU+OS memory pool can stall the whole box instead of raising a catchable OutOfMemoryError. NVIDIA Spark-class parts (DGX Spark / GB10, N1X "RTX Spark") have the same pool topology and the same failure mode, but only had an opt-in cap (UNSLOTH_SPARK_MEM_FRACTION). - worker.py: new _nvidia_classify_spark_unified_memory(props) mirroring _rocm_classify_unified_memory: is_integrated property first (authoritative on native Linux), then Spark device-name tokens -- WSL2's GPU paravirtualization masks is_integrated to 0 and renames the device (the N1X reports "JMJWOA-Generic-GPU"; verified on hardware), so the property alone misses Spark-under-WSL. Section 1h applies the 0.80 cap on match; UNSLOTH_SPARK_MEM_FRACTION overrides it and any value outside (0, 1] disables the guard. Discrete NVIDIA GPUs and CPU-only hosts are untouched. The existing generic OOM handler in the training loop surfaces the resulting OutOfMemoryError. - _utils.py: range-validate the opt-in UNSLOTH_SPARK_MEM_FRACTION -- "0" previously called set_per_process_memory_fraction(0.0), which makes every subsequent CUDA allocation OOM. - tests: test_spark_oom_guard.py mirroring test_rocm_oom_guard.py (property path, WSL name-token path, discrete negatives). 47/47 pass alongside the ROCm suite. Live-verified on the N1X (WSL2): classifier matches via JMJWOA, and with the cap set an over-allocation raises catchable torch.OutOfMemoryError instead of stalling the box; allocations recover after the error. Co-Authored-By: Claude Opus 4.8 --- studio/backend/core/training/worker.py | 72 +++++++++++++++ studio/backend/tests/test_spark_oom_guard.py | 95 ++++++++++++++++++++ unsloth/models/_utils.py | 6 +- 3 files changed, 172 insertions(+), 1 deletion(-) create mode 100644 studio/backend/tests/test_spark_oom_guard.py diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 50323478b2..0390e6745e 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -703,6 +703,36 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: return gcn_arch, is_unified +def _nvidia_classify_spark_unified_memory(props: Any) -> tuple[str, bool]: + """Classify an NVIDIA device as Spark-class unified-memory or discrete. + + Returns ``(marker, is_unified)``: + - ``marker``: the signal that matched (``"is_integrated"`` or the matching + device-name token), else ``""``. + - ``is_unified``: ``True`` for Spark-class parts that share one memory pool + with the OS (DGX Spark / GB10, N1X "RTX Spark", Grace-Blackwell desksides) + — these need the same lower ``set_per_process_memory_fraction`` cap as the + ROCm APUs: exhausting the shared pool can stall the whole box instead of + raising a catchable OutOfMemoryError. + + Classification priority: + 1. ``is_integrated`` device property (authoritative on native Linux). + 2. Device-name token match — WSL2's GPU paravirtualization masks + ``is_integrated`` to 0 and renames the device (the N1X reports + ``JMJWOA-Generic-GPU`` with ``is_integrated == 0``, verified on + hardware), so the property alone misses Spark-under-WSL. Tokens mirror + ``_DGX_SPARK_DEVICE_TOKENS`` in ``unsloth/models/_utils.py`` (duplicated + because this guard runs before any ML import). + """ + if getattr(props, "is_integrated", 0): + return "is_integrated", True + name_upper = (getattr(props, "name", "") or "").upper() + for token in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK"): + if token in name_upper: + return token, True + return "", False + + def _tilelang_platform_supported() -> bool: """True iff a tilelang 0.1.8 wheel will load: Linux x86_64/aarch64, non-HIP torch. @@ -2194,6 +2224,48 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> except Exception as _oom_guard_err: logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err) + # ── 1h. NVIDIA Spark-class unified-memory OOM guard ── + # Same failure mode as the ROCm APU guard above, NVIDIA flavor: Spark-class + # parts (DGX Spark / GB10, N1X "RTX Spark") share one memory pool with the + # OS, so over-allocation can stall the whole box instead of raising a + # catchable OutOfMemoryError. Cap the allocator at 0.80 like Strix Halo — + # the pool is shared with the host OS and page cache, so 20% headroom stays + # with the system. UNSLOTH_SPARK_MEM_FRACTION overrides the cap; any value + # outside (0, 1] disables the guard. Discrete NVIDIA GPUs are untouched + # (they already raise a graceful OOM). The generic OOM handler in the + # training loop surfaces the resulting OutOfMemoryError with remediation. + else: + try: + import torch as _torch_mem + if _torch_mem.cuda.is_available(): + _props = _torch_mem.cuda.get_device_properties(0) + _marker, _is_spark_uma = _nvidia_classify_spark_unified_memory(_props) + if _is_spark_uma: + _mem_fraction = 0.80 + _frac_env = os.environ.get("UNSLOTH_SPARK_MEM_FRACTION") + if _frac_env: + try: + _mem_fraction = float(_frac_env) + except ValueError: + _mem_fraction = 0.80 + if 0.0 < _mem_fraction <= 1.0: + _torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction) + logger.info( + "Spark unified-memory OOM guard: " + "set_per_process_memory_fraction(%.2f) — %s (matched %s)", + _mem_fraction, + _props.name, + _marker, + ) + else: + logger.info( + "Spark unified-memory OOM guard disabled " + "(UNSLOTH_SPARK_MEM_FRACTION=%s)", + _frac_env, + ) + except Exception as _oom_guard_err: + logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err) + # ── 2. Now import ML libraries (fresh in this clean process) ── try: _send_status(event_queue, "Importing Unsloth...") diff --git a/studio/backend/tests/test_spark_oom_guard.py b/studio/backend/tests/test_spark_oom_guard.py new file mode 100644 index 0000000000..74610d4369 --- /dev/null +++ b/studio/backend/tests/test_spark_oom_guard.py @@ -0,0 +1,95 @@ +# 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 _nvidia_classify_spark_unified_memory (Spark OOM-guard classifier). + +Two paths: (1) the ``is_integrated`` device property (authoritative on native +Linux), (2) device-name token match — needed because WSL2's GPU +paravirtualization masks ``is_integrated`` to 0 and renames the device (the N1X +reports ``JMJWOA-Generic-GPU``; verified on hardware). + +Mirrors test_rocm_oom_guard.py for the ROCm/Strix-Halo classifier the NVIDIA +guard was modeled on. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from core.training.worker import _nvidia_classify_spark_unified_memory + + +def _props(**kwargs) -> SimpleNamespace: + """Build a fake device-properties object with the given attributes.""" + return SimpleNamespace(**kwargs) + + +# ── Path 1: is_integrated property ─────────────────────────────────────────── + + +class TestIsIntegratedProperty: + """``is_integrated`` truthy means unified memory, regardless of name.""" + + def test_integrated_native_spark(self) -> None: + props = _props(is_integrated = 1, name = "NVIDIA GB10") + marker, is_unified = _nvidia_classify_spark_unified_memory(props) + assert marker == "is_integrated" + assert is_unified is True + + def test_integrated_wins_even_with_unknown_name(self) -> None: + props = _props(is_integrated = 1, name = "Some Future Unified Part") + marker, is_unified = _nvidia_classify_spark_unified_memory(props) + assert marker == "is_integrated" + assert is_unified is True + + +# ── Path 2: device-name token fallback (WSL masks is_integrated) ──────────── + + +class TestDeviceNameTokenFallback: + """is_integrated == 0 (or absent) -> classify by Spark name tokens.""" + + @pytest.mark.parametrize( + "name, expected_marker", + [ + ("JMJWOA-Generic-GPU", "JMJWOA"), # N1X under WSL2 (verified live) + ("NVIDIA GB10", "GB10"), # native DGX Spark + ("NVIDIA GB110", "GB110"), # "GB10" is not a substring of "GB110" + ("NVIDIA DGX Spark", "DGX SPARK"), + ("nvidia n1x prototype", "N1X"), # case-insensitive + ], + ) + def test_spark_names_unified(self, name: str, expected_marker: str) -> None: + props = _props(is_integrated = 0, name = name) + marker, is_unified = _nvidia_classify_spark_unified_memory(props) + assert is_unified is True + assert marker == expected_marker + + @pytest.mark.parametrize( + "name", + [ + "NVIDIA GeForce RTX 4090", + "NVIDIA H100 80GB HBM3", + "NVIDIA RTX 6000 Ada Generation", + "Tesla T4", + ], + ) + def test_discrete_names_not_unified(self, name: str) -> None: + props = _props(is_integrated = 0, name = name) + marker, is_unified = _nvidia_classify_spark_unified_memory(props) + assert is_unified is False + assert marker == "" + + def test_missing_attrs_defaults_discrete(self) -> None: + """No is_integrated, no name -> discrete (guard stays off).""" + marker, is_unified = _nvidia_classify_spark_unified_memory(_props()) + assert is_unified is False + assert marker == "" + + def test_none_name_defaults_discrete(self) -> None: + props = _props(is_integrated = 0, name = None) + marker, is_unified = _nvidia_classify_spark_unified_memory(props) + assert is_unified is False + assert marker == "" diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 270eb14fd6..c9e77336e0 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1049,7 +1049,11 @@ def patch_dgx_spark_runtime_defaults(): _frac = os.environ.get("UNSLOTH_SPARK_MEM_FRACTION") if _frac: try: - torch.cuda.set_per_process_memory_fraction(float(_frac)) + # Only (0, 1] is a usable cap: 0 would make EVERY allocation OOM + # and values > 1 are rejected by torch. Out-of-range = no cap. + _frac_val = float(_frac) + if 0.0 < _frac_val <= 1.0: + torch.cuda.set_per_process_memory_fraction(_frac_val) except Exception: pass