Fixed issue on DGX spark that caused unnessary moves of models causing slow down on unified memory

toolkit/util/unified_memory.py

@@ -0,0 +1,64 @@
+"""Unified memory (integrated GPU) support, e.g. NVIDIA DGX Spark (GB10).
+
+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.
+
+Set AITK_DISABLE_UNIFIED_MEMORY=1 to turn this off.
+"""
+import os
+import sys
+
+import torch
+
+_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"):
+        return False
+    # NVIDIA only. ROCm APUs report integrated too, but may have a dedicated VRAM carveout
+    if torch.version.hip is not None:
+        return False
+    try:
+        if not torch.cuda.is_available():
+            return False
+        return bool(torch.cuda.get_device_properties(0).is_integrated)
+    except Exception:
+        return False
+
+
+def apply_unified_memory_patches():
+    global _patched
+    if _patched or not is_unified_memory():
+        return
+    _patched = True
+
+    orig_to = torch.nn.Module.to
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, memory_format = torch._C._nn._parse_to(*args, **kwargs)
+        if device is None or device.type != "cpu":
+            return orig_to(self, *args, **kwargs)
+        # drop the cpu move, keep any dtype / memory format conversion
+        new_kwargs = {}
+        if dtype is not None:
+            new_kwargs["dtype"] = dtype
+        if memory_format is not None:
+            new_kwargs["memory_format"] = memory_format
+        if not new_kwargs:
+            return self
+        return orig_to(self, non_blocking=non_blocking, **new_kwargs)
+
+    def cpu(self):
+        return self
+
+    torch.nn.Module.to = to
+    torch.nn.Module.cpu = cpu
+    print(f"Unified memory device detected ({torch.cuda.get_device_name(0)}): module moves to cpu are disabled")