Register native pinned host memory with CUDA for GPU DMA

accelerator/abstract_accelerator.py

@@ -275,6 +275,14 @@ def _torch_pin_memory(self, tensor):
     def _torch_is_pinned(self, tensor):
         return tensor.is_pinned()
 
+    def register_host_memory(self, address, num_bytes):
+        """Register page-locked host memory with the active device runtime."""
+        return False
+
+    def unregister_host_memory(self, address):
+        """Unregister host memory previously registered with the device runtime."""
+        return None
+
     def pin_memory(self, tensor, make_copy=True, match_shape=True):
         from deepspeed.utils.pin_memory_tracker import track_pinned_memory
         track_pinned_memory(tensor.nbytes)

accelerator/cuda_accelerator.py

@@ -333,6 +333,15 @@ def on_accelerator(self, tensor):
         else:
             return False
 
+    def register_host_memory(self, address, num_bytes):
+        result = int(torch.cuda.cudart().cudaHostRegister(address, num_bytes, 0))
+        torch.cuda.check_error(result)
+        return True
+
+    def unregister_host_memory(self, address):
+        result = int(torch.cuda.cudart().cudaHostUnregister(address))
+        torch.cuda.check_error(result)
+
     def op_builder_dir(self):
         try:
             # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed

benchmarks/pin_memory/h2d_d2h_bench.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Compare torch and native pinned-memory H2D/D2H bandwidth on one CUDA GPU."""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+
+ARMS = {
+    "torch": {
+        "DS_PIN_MEMORY_BACKEND": "torch",
+        "DS_PIN_MEMORY_REGISTER_DEVICE": "1"
+    },
+    "native-unregistered": {
+        "DS_PIN_MEMORY_BACKEND": "native",
+        "DS_PIN_MEMORY_REGISTER_DEVICE": "0"
+    },
+    "native-registered": {
+        "DS_PIN_MEMORY_BACKEND": "native",
+        "DS_PIN_MEMORY_REGISTER_DEVICE": "1"
+    },
+}
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--arm", choices=ARMS)
+    parser.add_argument("--sizes-mib", type=int, nargs="+", default=[4, 64, 256])
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=50)
+    return parser.parse_args()
+
+
+def _time_copy(accelerator, copy_fn, stream, warmup, iters):
+    with accelerator.stream(stream):
+        for _ in range(warmup):
+            copy_fn()
+        stream.synchronize()
+        start = accelerator.Event(enable_timing=True)
+        end = accelerator.Event(enable_timing=True)
+        start.record(stream)
+        for _ in range(iters):
+            copy_fn()
+        end.record(stream)
+    stream.synchronize()
+    return start.elapsed_time(end) / 1000.0 / iters
+
+
+def _allocate_host(accelerator, numel, arm):
+    if arm == "torch":
+        return torch.empty(numel, dtype=torch.float32, pin_memory=True)
+
+    return accelerator.pin_memory(torch.empty(numel, dtype=torch.float32), make_copy=False)
+
+
+def _run_arm(args):
+    for key, value in ARMS[args.arm].items():
+        os.environ[key] = value
+
+    accelerator = get_accelerator()
+    if accelerator.device_name() != "cuda" or not accelerator.is_available():
+        raise RuntimeError("CUDA GPU is required")
+    accelerator.set_device(0)
+    stream = accelerator.Stream()
+
+    for size_mib in args.sizes_mib:
+        num_bytes = size_mib * 1024 * 1024
+        numel = num_bytes // torch.tensor([], dtype=torch.float32).element_size()
+        host = _allocate_host(accelerator, numel, args.arm)
+        device = torch.empty_like(host, device=accelerator.current_device_name())
+
+        h2d_seconds = _time_copy(accelerator, lambda: device.copy_(host, non_blocking=True), stream, args.warmup,
+                                 args.iters)
+        d2h_seconds = _time_copy(accelerator, lambda: host.copy_(device, non_blocking=True), stream, args.warmup,
+                                 args.iters)
+
+        result = {
+            "arm": args.arm,
+            "size_mib": size_mib,
+            "h2d_gbps": num_bytes / h2d_seconds / 1e9,
+            "d2h_gbps": num_bytes / d2h_seconds / 1e9,
+            "torch_is_pinned": host.is_pinned(),
+            "accelerator_is_pinned": accelerator.is_pinned(host),
+        }
+        print(f"RESULT={json.dumps(result, sort_keys=True)}", flush=True)
+        accelerator.unpin_memory(host)
+
+
+def _run_all(args):
+    results = []
+    for arm in ARMS:
+        command = [
+            sys.executable,
+            os.path.abspath(__file__),
+            "--arm",
+            arm,
+            "--sizes-mib",
+            *(str(size) for size in args.sizes_mib),
+            "--warmup",
+            str(args.warmup),
+            "--iters",
+            str(args.iters),
+        ]
+        process = subprocess.run(command, check=True, text=True, capture_output=True)
+        if process.stderr:
+            print(process.stderr, file=sys.stderr, end="")
+        for line in process.stdout.splitlines():
+            print(line)
+            if line.startswith("RESULT="):
+                results.append(json.loads(line.removeprefix("RESULT=")))
+
+    print("\narm,size_mib,h2d_gbps,d2h_gbps,torch_is_pinned,accelerator_is_pinned")
+    for result in results:
+        print(f"{result['arm']},{result['size_mib']},{result['h2d_gbps']:.2f},{result['d2h_gbps']:.2f},"
+              f"{result['torch_is_pinned']},{result['accelerator_is_pinned']}")
+
+
+if __name__ == "__main__":
+    arguments = _parse_args()
+    if arguments.arm:
+        _run_arm(arguments)
+    else:
+        _run_all(arguments)

deepspeed/utils/pin_memory.py

@@ -5,6 +5,8 @@
 import os
 import weakref
 
+from deepspeed.utils import logger
+
 # ``torch._subclasses.fake_tensor`` is a private API that may be absent on some
 # torch versions; guard the import so this module stays importable (including
 # during setup when torch may not be installed).
@@ -25,6 +27,8 @@ def __init__(self):
         # base address -> weakref.finalize handle for the returned tensor, so an
         # explicit unpin() can cancel the GC-triggered free.
         self._finalizers = {}
+        # Allocation bases successfully registered with the active device runtime.
+        self._device_registered = set()
         # Fail early: native pinning is useless without the pin_memory handle, so
         # surface the build/load failure here instead of silently degrading.
         try:
@@ -43,6 +47,14 @@ def pin(self, tensor, make_copy=True, match_shape=True):
         base = self._handle.new_cpu_locked_tensor(numel, tensor)
         begin = base.data_ptr()
         locked = base[:numel]
+        if base.nbytes and self._device_registration_enabled():
+            from deepspeed.accelerator import get_accelerator
+            try:
+                if get_accelerator().register_host_memory(begin, base.nbytes):
+                    self._device_registered.add(begin)
+            except Exception as e:
+                logger.warning_once(
+                    f"Native pinned-memory device registration failed; continuing with mlock only: {e}")
         if make_copy:
             locked.copy_(tensor.reshape(-1))
         if match_shape:
@@ -59,7 +71,7 @@ def pin(self, tensor, make_copy=True, match_shape=True):
         # (not the returned view) and frees by address; ``base`` must not be passed
         # as a finalize argument or it would be kept alive forever.
         self._finalizers[begin] = weakref.finalize(base, self._release, self._handle, begin, self._ranges,
-                                                   self._finalizers)
+                                                   self._finalizers, self._device_registered)
         return locked
 
     def is_pinned(self, tensor):
@@ -77,6 +89,9 @@ def unpin(self, tensor):
         begin = getattr(tensor, "ds_pin_base", None)
         if begin is None:
             begin = tensor.data_ptr()
+        # Unregister first. If this fails, keep the allocation so the driver is
+        # not left holding a registration for pages later reused by malloc.
+        self._unregister_device(begin, self._device_registered)
         finalizer = self._finalizers.pop(begin, None)
         if finalizer is not None:
             # Explicit unpin owns the free; cancel the GC finalizer to avoid a
@@ -89,7 +104,13 @@ def unpin(self, tensor):
         return freed
 
     @staticmethod
-    def _release(handle, begin, ranges, finalizers):
+    def _release(handle, begin, ranges, finalizers, device_registered):
+        try:
+            NativePinnedMemory._unregister_device(begin, device_registered)
+        except Exception:
+            # Interpreter shutdown / dead device context: leave the allocation
+            # so the driver does not retain a registration for recycled pages.
+            return
         ranges.pop(begin, None)
         finalizers.pop(begin, None)
         try:
@@ -99,6 +120,23 @@ def _release(handle, begin, ranges, finalizers):
             # during interpreter shutdown.
             pass
 
+    @staticmethod
+    def _unregister_device(begin, device_registered):
+        if begin not in device_registered:
+            return
+        from deepspeed.accelerator import get_accelerator
+        get_accelerator().unregister_host_memory(begin)
+        device_registered.discard(begin)
+
+    @staticmethod
+    def _device_registration_enabled():
+        value = os.environ.get("DS_PIN_MEMORY_REGISTER_DEVICE", "1").strip().lower()
+        if value in ("1", "true", "yes", "on"):
+            return True
+        if value in ("0", "false", "no", "off"):
+            return False
+        raise ValueError("DS_PIN_MEMORY_REGISTER_DEVICE must be one of: 1, 0, true, false, yes, no, on, off")
+
     @staticmethod
     def _has_real_storage(tensor):
         # Fake/meta tensors have no storage; skip the range check to avoid a

docs/code-docs/source/memory.rst

@@ -356,7 +356,8 @@ counted by ``track_pinned_memory`` when pages are actually locked. Differences:
      - ``native``
    * - Allocator
      - ``torch.Tensor.pin_memory()`` (device-specific accelerator hook)
-     - DeepNVMe page-locked allocator (``posix_memalign`` + ``mlock``) via pin_memory
+     - DeepNVMe page-locked allocator (``posix_memalign`` + ``mlock``) via pin_memory;
+       registered with the CUDA runtime by default for asynchronous DMA
    * - Selection
      - ``DS_PIN_MEMORY_BACKEND`` unset or ``torch``
      - ``DS_PIN_MEMORY_BACKEND=native``
@@ -386,6 +387,25 @@ Example:
     export DS_PIN_MEMORY_BACKEND=native
     deepspeed train.py ...
 
+Native device registration
+==========================
+
+Native allocations are device-independent ``mlock`` buffers. On CUDA systems,
+DeepSpeed additionally calls ``cudaHostRegister`` so PyTorch can use them for
+asynchronous H2D/D2H DMA. Device registration is enabled by default and can be
+disabled for comparison or debugging:
+
+.. code-block:: bash
+
+    export DS_PIN_MEMORY_BACKEND=native
+    export DS_PIN_MEMORY_REGISTER_DEVICE=0  # mlock only; default is 1
+
+``DS_PIN_MEMORY_REGISTER_DEVICE`` accepts ``1``/``0``, ``true``/``false``,
+``yes``/``no``, and ``on``/``off``. Accelerators without a registration hook
+continue to use the device-independent ``mlock`` buffer. If registration fails,
+DeepSpeed logs a warning and retains the valid ``mlock`` allocation for CPU and
+AIO use.
+
 Requirements for native
 =======================