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 <noreply@anthropic.com>

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...")

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 == ""

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