diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py index b94ff56dec..c5cd537b89 100644 --- a/unsloth/kernels/flex_attention.py +++ b/unsloth/kernels/flex_attention.py @@ -25,6 +25,40 @@ torch_compile_options = { "triton.cudagraphs": False, } + +def _flex_is_dgx_spark(): + # Mirror of unsloth.models._utils.is_dgx_spark(), inlined to avoid importing + # `unsloth.models` from this low-level `kernels` module (circular at import). + # DGX Spark / N1X = aarch64 + NVIDIA CUDA + a Spark device-name token. + _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( + t in names for t in ("GB10", "JMJWOA", "N1X", "DGX SPARK", "GB110") + ) + except Exception: + return False + + +# DGX Spark / N1X has 48 SMs (< inductor's 68-SM is_big_gpu threshold), so +# max_autotune_gemm is already skipped; dropping max_autotune only saves the +# wasted compile-time search -- identical kernels, no accuracy/throughput change. +if _flex_is_dgx_spark(): + torch_compile_options["max_autotune"] = False + # Flex Attention supported from torch 2.5 onwards only try: from torch.nn.attention.flex_attention import ( diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index c9abbce183..8965cac7e0 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1034,8 +1034,74 @@ def patch_dgx_spark_memory_config(): ) + "expandable_segments:True" +def patch_dgx_spark_runtime_defaults(): + """Spark UMA runtime defaults (accuracy-neutral, gated, env-overridable). + + - `UNSLOTH_DISABLE_DOUBLE_BUFFER=1`: unsloth-zoo's gradient-checkpointing + double-buffer is enabled via a `torch.cuda.mem_get_info` free-memory check + that UNDERCOUNTS on UMA, and it stages an extra GPU buffer to overlap a + host<->device copy that is physically free on a shared pool. Default it off + on Spark (`setdefault`, so a user can still force it back on). Must be set + before unsloth-zoo initializes gradient checkpointing -- `import unsloth` + precedes that, so this is in time. + - `set_per_process_memory_fraction`: OPT-IN safety valve. On Spark UMA an + over-allocation can wedge the box (untracked UMA allocations may never trip + a catchable OOM). If the user sets `UNSLOTH_SPARK_MEM_FRACTION=<0..1>`, cap + the caching allocator so it raises OutOfMemoryError early. Default unset -> + NO cap (no capacity loss); purely opt-in. + Strict no-op off-Spark. + """ + if not is_dgx_spark(): + return + os.environ.setdefault("UNSLOTH_DISABLE_DOUBLE_BUFFER", "1") + _frac = os.environ.get("UNSLOTH_SPARK_MEM_FRACTION") + if _frac: + try: + torch.cuda.set_per_process_memory_fraction(float(_frac)) + except Exception: + pass + + +def patch_dgx_spark_dataloader_defaults(): + """On Spark UMA, default `dataloader_pin_memory` to False (accuracy-neutral). + + Page-locked host memory exists to speed host->device DMA; on unified memory + there is no separate device memory, so pinning only reserves non-pageable RAM + from the shared pool and adds a staging copy -- pure waste. Mirrors + transformers' own `if self.use_cpu: self.dataloader_pin_memory = False` + precedent. Wraps the base `TrainingArguments.__post_init__`, so SFT + every + TRL trainer (whose configs call `super().__post_init__()`) are covered with + one idempotent patch. Only flips the library default `True`; opt out with + `UNSLOTH_SPARK_KEEP_PIN_MEMORY=1`. Strict no-op off-Spark; never changes any + computed value, so accuracy is unaffected. + """ + if not is_dgx_spark(): + return + if os.environ.get("UNSLOTH_SPARK_KEEP_PIN_MEMORY") == "1": + return + try: + from transformers import training_args as _ta + + Base = _ta.TrainingArguments + except Exception: + return + if getattr(Base.__post_init__, "_unsloth_spark_uma", False): + return + _orig_post_init = Base.__post_init__ + + def __post_init__(self): + _orig_post_init(self) + if getattr(self, "dataloader_pin_memory", None) is True: + self.dataloader_pin_memory = False + + __post_init__._unsloth_spark_uma = True + Base.__post_init__ = __post_init__ + + patch_dgx_spark_memory_config() patch_dgx_spark_caching_allocator_warmup() +patch_dgx_spark_runtime_defaults() +patch_dgx_spark_dataloader_defaults() class _RaiseUninitialized(logging.Handler): @@ -1620,6 +1686,14 @@ torch_compile_options = { "trace.enabled": UNSLOTH_COMPILE_DEBUG, "triton.cudagraphs": False, } +# DGX Spark / N1X: this GPU has 48 SMs, below inductor's hardcoded 68-SM +# `is_big_gpu` threshold, so `max_autotune_gemm` is already skipped by inductor +# (the "Not enough SMs to use max_autotune_gemm mode" warning). Dropping +# max_autotune on Spark only avoids the wasted compile-time autotuning search -- +# the produced Triton/inductor kernels are identical, so steady-state throughput +# and accuracy are unchanged. Strict no-op off-Spark (gated by is_dgx_spark()). +if is_dgx_spark(): + torch_compile_options["max_autotune"] = False import accelerate