diff --git a/accelerator/abstract_accelerator.py b/accelerator/abstract_accelerator.py index d03699a1c922..718cf872b740 100644 --- a/accelerator/abstract_accelerator.py +++ b/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) diff --git a/accelerator/cuda_accelerator.py b/accelerator/cuda_accelerator.py index 1fc6f88aaba3..aab493d65f41 100644 --- a/accelerator/cuda_accelerator.py +++ b/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 diff --git a/benchmarks/pin_memory/h2d_d2h_bench.py b/benchmarks/pin_memory/h2d_d2h_bench.py new file mode 100644 index 000000000000..0cad917cdbb9 --- /dev/null +++ b/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) diff --git a/deepspeed/utils/pin_memory.py b/deepspeed/utils/pin_memory.py index 1aecfe484a5d..1bcda8ec01b6 100644 --- a/deepspeed/utils/pin_memory.py +++ b/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 diff --git a/docs/code-docs/source/memory.rst b/docs/code-docs/source/memory.rst index f00db08621ab..c59f2936b1ad 100644 --- a/docs/code-docs/source/memory.rst +++ b/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 ======================= diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py index 09875dae51c2..ec2e924d0ef3 100644 --- a/tests/unit/v1/pin_memory/test_pin_memory.py +++ b/tests/unit/v1/pin_memory/test_pin_memory.py @@ -5,6 +5,8 @@ import pytest import torch +from deepspeed.accelerator.cpu_accelerator import CPU_Accelerator +from deepspeed.accelerator.cuda_accelerator import CUDA_Accelerator from deepspeed.utils.pin_memory import NativePinnedMemory @@ -114,3 +116,162 @@ def test_is_pinned_handles_storageless_tensors(native_pins): meta_tensor = torch.zeros(2, 8, device="meta") assert native_pins.is_pinned(meta_tensor) is False + + +class _RegisteringAccelerator: + + def __init__(self): + self.registered = [] + self.unregistered = [] + + def register_host_memory(self, address, num_bytes): + self.registered.append((address, num_bytes)) + return True + + def unregister_host_memory(self, address): + self.unregistered.append(address) + + +def test_native_device_registration_and_unpin(monkeypatch, native_pins): + accelerator = _RegisteringAccelerator() + monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: accelerator) + monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1") + + pinned = native_pins.pin(torch.empty(32), make_copy=False) + begin = pinned.data_ptr() + assert accelerator.registered == [(begin, pinned.nbytes)] + assert begin in native_pins._device_registered + + assert native_pins.unpin(pinned) is True + assert accelerator.unregistered == [begin] + assert begin not in native_pins._device_registered + # A second call must not unregister or free the allocation twice. + assert native_pins.unpin(pinned) is False + assert accelerator.unregistered == [begin] + + +def test_native_device_registration_disabled(monkeypatch, native_pins): + accelerator = _RegisteringAccelerator() + monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: accelerator) + monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "0") + + pinned = native_pins.pin(torch.empty(32), make_copy=False) + assert native_pins.is_pinned(pinned) + assert accelerator.registered == [] + assert native_pins.unpin(pinned) is True + assert accelerator.unregistered == [] + + +def test_cpu_device_registration_is_noop(): + accelerator = CPU_Accelerator() + assert accelerator.register_host_memory(1234, 4096) is False + assert accelerator.unregister_host_memory(1234) is None + + +def test_cuda_device_registration_calls_cudart(monkeypatch): + + class _Cudart: + + def __init__(self): + self.registered = [] + self.unregistered = [] + + def cudaHostRegister(self, address, num_bytes, flags): + self.registered.append((address, num_bytes, flags)) + return 0 + + def cudaHostUnregister(self, address): + self.unregistered.append(address) + return 0 + + cudart = _Cudart() + errors = [] + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) #ignore-cuda + monkeypatch.setattr(torch.cuda, "check_error", errors.append) #ignore-cuda + accelerator = CUDA_Accelerator.__new__(CUDA_Accelerator) + + assert accelerator.register_host_memory(1234, 4096) is True + accelerator.unregister_host_memory(1234) + assert cudart.registered == [(1234, 4096, 0)] + assert cudart.unregistered == [1234] + assert errors == [0, 0] + + +def test_cpu_native_pin_with_register_env_on(monkeypatch, native_pins): + """CPU accelerator has no register hook; native pin still works with default-on.""" + monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1") + monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: CPU_Accelerator()) + pinned = native_pins.pin(torch.empty(32), make_copy=False) + assert native_pins.is_pinned(pinned) + assert native_pins.unpin(pinned) is True + + +def test_device_registration_failure_keeps_mlock(monkeypatch, native_pins): + + class _FailingAccelerator: + + def register_host_memory(self, address, num_bytes): + raise RuntimeError("simulated cudaHostRegister failure") + + def unregister_host_memory(self, address): + raise AssertionError("unregister must not run when register failed") + + monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: _FailingAccelerator()) + monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1") + pinned = native_pins.pin(torch.empty(32), make_copy=False) + assert native_pins.is_pinned(pinned) + assert pinned.data_ptr() not in native_pins._device_registered + assert native_pins.unpin(pinned) is True + + +def test_device_registration_gc_unregisters(monkeypatch, native_pins): + import gc + + accelerator = _RegisteringAccelerator() + monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: accelerator) + monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1") + pinned = native_pins.pin(torch.empty(32), make_copy=False) + begin = pinned.data_ptr() + del pinned + gc.collect() + assert begin not in native_pins._ranges + assert accelerator.unregistered == [begin] + + +def test_invalid_register_device_env(monkeypatch, native_pins): + monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "maybe") + with pytest.raises(ValueError, match="DS_PIN_MEMORY_REGISTER_DEVICE"): + native_pins.pin(torch.empty(8), make_copy=False) + + +def test_unpin_keeps_allocation_when_unregister_fails(monkeypatch, native_pins): + + class _UnregisterFailAccelerator(_RegisteringAccelerator): + + def __init__(self): + super().__init__() + self.fail_unregister = True + + def unregister_host_memory(self, address): + if self.fail_unregister: + raise RuntimeError("simulated cudaHostUnregister failure") + super().unregister_host_memory(address) + + accelerator = _UnregisterFailAccelerator() + monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: accelerator) + monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1") + pinned = native_pins.pin(torch.empty(32), make_copy=False) + begin = pinned.data_ptr() + + with pytest.raises(RuntimeError, match="cudaHostUnregister"): + native_pins.unpin(pinned) + + assert native_pins.is_pinned(pinned) + assert begin in native_pins._device_registered + assert begin in native_pins._ranges + assert accelerator.unregistered == [] + + accelerator.fail_unregister = False + assert native_pins.unpin(pinned) is True + assert accelerator.unregistered == [begin] + assert begin not in native_pins._device_registered