From 221cac39a02890b022386933cd96e489cf8b04e5 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 27 Oct 2026 08:13:02 -0600 Subject: [PATCH] Improve unified memory usage on the RTX Spark --- toolkit/util/unified_memory.py | 22 ++++++++++++------- ui/src/server/monitor.ts | 40 ++++++++++++++++++++++++---------- 2 files changed, 43 insertions(+), 19 deletions(-) diff --git a/toolkit/util/unified_memory.py b/toolkit/util/unified_memory.py index 4229e32..82291b7 100644 --- a/toolkit/util/unified_memory.py +++ b/toolkit/util/unified_memory.py @@ -1,13 +1,19 @@ -"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10). +"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10) and +the Windows RTX Spark (N1X). On these devices the CPU and GPU share the same physical memory, so moving a module to the CPU frees nothing. Worse, device -> host copies are extremely -slow (~100MB/s on GB10), so every "offload to cpu" of a large transformer or -text encoder stalls for minutes. When a unified memory device is detected we -make module level moves to the CPU a no-op (dtype casts still apply). Tensor -level .cpu()/.to('cpu') is left alone since saving, numpy, etc. need it. +slow (~100MB/s on GB10 under Linux, ~0.5-2GB/s pageable on Windows/WDDM), so +every "offload to cpu" of a large transformer or text encoder stalls. When a +unified memory device is detected we make module level moves to the CPU a +no-op (dtype casts still apply). Tensor level .cpu()/.to('cpu') is left alone +since saving, numpy, etc. need it. -Set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off. +Note for Windows: WDDM gives CUDA a fixed budget (~80% of system RAM, e.g. +~45GB of 54GB) that host allocations do not count against. nvidia-smi / Task +Manager only show the small dedicated segment (8GB); the real pool is what +torch.cuda.mem_get_info() reports. If a job needs to keep modules on the CPU +to stay under that budget, set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off. """ import os import sys @@ -20,8 +26,8 @@ _patched = False def is_unified_memory() -> bool: if os.environ.get("AITK_DISABLE_UNIFIED_MEMORY", "0") == "1": return False - # linux only (DGX Spark). The Windows RTX Spark stack does not behave the same way - if not sys.platform.startswith("linux"): + # macOS has no CUDA; MPS unified memory is handled elsewhere + if sys.platform == "darwin": return False # NVIDIA only. ROCm APUs report integrated too, but may have a dedicated VRAM carveout if torch.version.hip is not None: diff --git a/ui/src/server/monitor.ts b/ui/src/server/monitor.ts index 31a34bf..23fa11e 100644 --- a/ui/src/server/monitor.ts +++ b/ui/src/server/monitor.ts @@ -36,11 +36,20 @@ const NV_BATCH_FLUSH_MS = 100; const CPU_TEMP_REFRESH_MS = 5000; // Power limits for GPUs whose driver reports none (nvidia-smi gives [N/A]). -// GB10 (DGX Spark) is a unified SoC with no settable GPU limit; 140 W is -// NVIDIA's rated TDP for the GB10 superchip. -const KNOWN_POWER_LIMITS_W: Record = { - 'NVIDIA GB10': 140, -}; +// These unified SoCs have no settable GPU limit, so show the rated TDP of the +// whole chip. GB10 (DGX Spark) is 140 W. RTX Spark N1X: the 5120-core part is +// laptop-only, rated 45-80 W (OEM configurable); the 6144-core part is 140 W +// on desktop and 45-80 W in laptops but shares its name, so it falls through +// to the desktop figure. +const KNOWN_POWER_LIMITS_W: [RegExp, number][] = [ + [/\bGB10\b/, 140], + [/RTX Spark.*5120-core/, 80], + [/RTX Spark/, 140], +]; + +function knownPowerLimit(name: string): number { + return KNOWN_POWER_LIMITS_W.find(([re]) => re.test(name))?.[1] ?? NaN; +} function parseGpuLine(line: string): GpuInfo | null { const [ @@ -76,7 +85,7 @@ function parseGpuLine(line: string): GpuInfo | null { }, power: { draw: parseFloat(powerDraw), - limit: isNaN(parseFloat(powerLimit)) ? (KNOWN_POWER_LIMITS_W[name] ?? NaN) : parseFloat(powerLimit), + limit: isNaN(parseFloat(powerLimit)) ? knownPowerLimit(name) : parseFloat(powerLimit), }, clocks: { graphics: parseInt(clockGraphics), @@ -88,19 +97,28 @@ function parseGpuLine(line: string): GpuInfo | null { }; } +// Integrated NVIDIA GPUs that share system RAM with the CPU. +const UNIFIED_MEMORY_GPU = /\b(GB10|RTX Spark)\b/i; + /** - * Unified-memory GPUs (GB10 / DGX Spark) have no dedicated VRAM, so - * nvidia-smi reports memory as [N/A]. Report the system RAM pool they share - * instead, like the Mac path does for Apple silicon. + * Unified-memory GPUs (GB10 / DGX Spark, RTX Spark) have no dedicated VRAM. + * On Linux nvidia-smi reports their memory as [N/A]; on Windows (WDDM) it + * reports only the small dedicated segment (8 GB) even though CUDA can use + * most of system RAM. Report the system RAM pool they share instead, like + * the Mac path does for Apple silicon. */ +function isSharedMemoryGpu(g: GpuInfo): boolean { + return isNaN(g.memory.total) || UNIFIED_MEMORY_GPU.test(g.name); +} + function withSharedMemory(gpu: GPUApiResponse, cpu: CpuInfo | null): GPUApiResponse { - if (!cpu || !gpu.gpus.some(g => isNaN(g.memory.total))) return gpu; + if (!cpu || !gpu.gpus.some(isSharedMemoryGpu)) return gpu; const total = Math.round(cpu.totalMemory); const free = Math.round(cpu.availableMemory); return { ...gpu, gpus: gpu.gpus.map(g => - isNaN(g.memory.total) ? { ...g, memory: { total, free, used: total - free, shared: true } } : g, + isSharedMemoryGpu(g) ? { ...g, memory: { total, free, used: total - free, shared: true } } : g, ), }; }