DGX Spark / N1X: shared is_dgx_spark() + caching_allocator_warmup no-op

Runtime training support for NVIDIA Blackwell unified-memory (UMA) machines --
DGX Spark (GB10) and the N1X "RTX Spark" laptop. Gated so it is a strict no-op
on every non-Spark platform (x86 NVIDIA, AMD/ROCm, Intel/XPU, Mac/MLX, discrete
aarch64 GH200/GB200): those are not aarch64 and/or report non-matching device
names, so behaviour there is byte-for-byte unchanged.

models/_utils.py:
- Add is_dgx_spark(): aarch64 + NVIDIA CUDA + a known Spark device-name token
  (GB10 / JMJWOA / N1X / ...). @lru_cache; overridable via UNSLOTH_FORCE_DGX_SPARK.
  One shared detector that also catches the N1X laptop, which reports
  "JMJWOA-Generic-GPU" rather than "NVIDIA GB10".
- Add patch_dgx_spark_caching_allocator_warmup(), applied at import: no-ops
  transformers.modeling_utils.caching_allocator_warmup on Spark. HF sizes a GPU
  pre-allocation from cudaMemGetInfo() to warm the caching allocator; on Spark
  UMA cudaMemGetInfo undercounts free memory (reclaimable buffer cache shows as
  unavailable), so the warmup torch.empty() raises
  `AcceleratorError: invalid argument` and aborts any bitsandbytes 4/8-bit load.
  The warmup is only a speed hint -> dropping it on Spark lets quantized loads
  succeed. Idempotent; single call site (modeling_utils.py:4212) confirmed.
  (Patch credited to Roland [UnAI] / Daniel, Unsloth Discord.)

models/loader.py:
- Replace the two inline `"NVIDIA GB10" in get_device_name()` checks (which
  disable the currently-broken vLLM fast_inference) with is_dgx_spark(), so the
  N1X is covered too. Same behaviour on DGX Spark; no change off-Spark.

torch.compile + Triton are verified WORKING on the N1X (Triton 3.6.0; a real
gemma-3-270m-it 4-bit finetune with compile ON trains and emits the full compiled
cache), so nothing is disabled -- UNSLOTH_COMPILE_DISABLE is not set.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

unsloth/models/_utils.py

@@ -939,6 +939,72 @@ except:
 from transformers.modeling_utils import logger as transformers_logger
 
 
+# ---- NVIDIA DGX Spark (GB10) / N1X "RTX Spark" (Blackwell unified-memory) support ----
+# These Blackwell unified-memory (UMA) machines report different device names:
+# "NVIDIA GB10" on DGX Spark, "JMJWOA-Generic-GPU" on the pre-launch N1X laptop.
+# One shared detector so every Spark-specific workaround uses the same definition.
+# The aarch64 + CUDA gate makes this a strict no-op on x86_64 NVIDIA, AMD/ROCm,
+# Intel/XPU, Mac/MLX, and discrete aarch64 GPUs (GH200/GB200) -- those report
+# non-matching names and/or are not aarch64, so behaviour there is unchanged.
+_DGX_SPARK_DEVICE_TOKENS = ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110")
+
+@functools.lru_cache(maxsize = None)
+def is_dgx_spark():
+    """True only on a DGX Spark / N1X Spark-class machine.
+
+    Gate: aarch64 + NVIDIA CUDA + a known Spark device-name token. Overridable for
+    testing via UNSLOTH_FORCE_DGX_SPARK=1 (force on) / =0 (force off).
+    """
+    _force = os.environ.get("UNSLOTH_FORCE_DGX_SPARK")
+    if _force == "1": return True
+    if _force == "0": return False
+    try:
+        import platform
+        if platform.machine().lower() not in ("aarch64", "arm64"):
+            return False
+        if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
+            return False
+        names = " ".join(
+            str(torch.cuda.get_device_name(i)).upper()
+            for i in range(torch.cuda.device_count())
+        )
+        return any(token in names for token in _DGX_SPARK_DEVICE_TOKENS)
+    except Exception:
+        return False
+pass
+
+
+def patch_dgx_spark_caching_allocator_warmup():
+    """No-op `transformers.modeling_utils.caching_allocator_warmup` on Spark UMA.
+
+    HF sizes a GPU pre-allocation from `cudaMemGetInfo()` to warm the caching
+    allocator. On Spark unified memory `cudaMemGetInfo` undercounts free memory
+    (reclaimable buffer cache is reported unavailable), so the warmup
+    `torch.empty(...)` raises `AcceleratorError: invalid argument` and aborts any
+    runtime-quantized (bitsandbytes 4/8-bit) load. The warmup is only a speed hint,
+    so skipping it on Spark merely forgoes a minor warmup while letting loads
+    succeed. No-op on every non-Spark platform (gated by `is_dgx_spark()`).
+    Idempotent: re-applying is a no-op (marked via `_unsloth_spark_noop`).
+    """
+    if not is_dgx_spark():
+        return
+    try:
+        from transformers import modeling_utils as _mu
+    except Exception:
+        return
+    if not hasattr(_mu, "caching_allocator_warmup"):
+        return
+    if getattr(_mu.caching_allocator_warmup, "_unsloth_spark_noop", False):
+        return
+    def _noop(*args, **kwargs):
+        return None
+    _noop._unsloth_spark_noop = True
+    _mu.caching_allocator_warmup = _noop
+pass
+
+patch_dgx_spark_caching_allocator_warmup()
+
+
 class _RaiseUninitialized(logging.Handler):
     def __init__(self):
         super().__init__()

unsloth/models/loader.py

@@ -16,6 +16,7 @@ from ._utils import (
     _prepare_model_for_qat,
     is_bfloat16_supported,
     is_vLLM_available,
+    is_dgx_spark,
     HAS_FLASH_ATTENTION,
     HAS_FLASH_ATTENTION_SOFTCAPPING,
     USE_MODELSCOPE,
@@ -379,10 +380,10 @@ class FastLanguageModel(FastLlamaModel):
                 )
             if DEVICE_TYPE_TORCH == "cuda":
                 for i in range(DEVICE_COUNT):
-                    # [TODO] DGX Spark vLLM breaks
-                    if "NVIDIA GB10" in str(torch.cuda.get_device_name(i)).upper():
+                    # [TODO] DGX Spark / N1X (Spark-class) vLLM breaks
+                    if is_dgx_spark():
                         print(
-                            "Unsloth: DGX Spark detected - `fast_inference=True` is currently broken as of January 2026.\n"
+                            "Unsloth: DGX Spark / N1X (Spark-class GPU) detected - `fast_inference=True` is currently broken as of January 2026.\n"
                             "Defaulting to native Unsloth inference."
                         )
                         fast_inference = False
@@ -1005,10 +1006,10 @@ class FastModel(FastBaseModel):
                 )
             if DEVICE_TYPE_TORCH == "cuda":
                 for i in range(DEVICE_COUNT):
-                    # [TODO] DGX Spark vLLM breaks
-                    if "NVIDIA GB10" in str(torch.cuda.get_device_name(i)).upper():
+                    # [TODO] DGX Spark / N1X (Spark-class) vLLM breaks
+                    if is_dgx_spark():
                         print(
-                            "Unsloth: DGX Spark detected - `fast_inference=True` is currently broken as of January 2026.\n"
+                            "Unsloth: DGX Spark / N1X (Spark-class GPU) detected - `fast_inference=True` is currently broken as of January 2026.\n"
                             "Defaulting to native Unsloth inference."
                         )
                         fast_inference = False