diff --git a/accelerator/abstract_accelerator.py b/accelerator/abstract_accelerator.py index 96a671e7f3c8..9f33b4f75bbc 100644 --- a/accelerator/abstract_accelerator.py +++ b/accelerator/abstract_accelerator.py @@ -286,6 +286,10 @@ def unregister_host_memory(self, address): """Unregister host memory previously registered with the device runtime.""" return None + def pin_memory_alignment(self): + """Byte alignment required for device-registered host memory; 1 means none.""" + return 1 + # CPU torch pinning is a historical no-op; subclasses that really page-lock # keep the default True so tracker accounting matches the docs. _torch_pins_host_memory = True diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py index fbcc10172ac8..2b9abfab9507 100644 --- a/accelerator/npu_accelerator.py +++ b/accelerator/npu_accelerator.py @@ -13,6 +13,33 @@ except ImportError: pass +ACL_SUCCESS = 0 +# The host-registration flag is an ACL_HOST_REG_* bitmask. MAPPED page-locks the +# range and adds a device mapping (which DeepSpeed never reads through); it is +# chosen over PINNED because measured H2D on this platform falls back to +# mlock-only speed (~8-9 GB/s) for PINNED-only registrations of larger buffers +# (64 MiB), while MAPPED keeps full DMA bandwidth (~23 GB/s). MAPPED requires +# 4K-aligned addresses, which the native allocator guarantees (posix_memalign); +# see the aclrtHostRegisterV2 API reference: +# https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/latest/API/runtimeapi/aclcppdevg_03_2128.html +ACL_HOST_REG_MAPPED = 0x2 + + +def _npu_host_copy_funcs(): + """Resolve the npurt host-registration functions, or (None, reason). + + torch.npu.npurt() returns the runtime-API module exposing + npuHostRegister/npuHostUnregister; the binding is maintained with + torch_npu and initializes the runtime itself. + """ + if not hasattr(torch, "npu") or not hasattr(torch.npu, "npurt"): + return None, "torch.npu.npurt is unavailable in this torch_npu build" + try: + npurt = torch.npu.npurt() + return (npurt.npuHostRegister, npurt.npuHostUnregister), None + except RuntimeError: + return None, "torch.npu.npurt() failed to initialize the NPU runtime" + class NPU_Accelerator(DeepSpeedAccelerator): @@ -152,6 +179,43 @@ def total_memory(self, device_index=None): def available_memory(self, device_index=None): return self.total_memory(device_index) - self.memory_allocated(device_index) + # Host memory registration + def pin_memory_alignment(self): + # MAPPED registration requires 4K-aligned addresses (per the + # aclrtHostRegisterV2 API reference cited on ACL_HOST_REG_MAPPED). + # NativePinnedMemory rounds the range down to this alignment before + # calling register_host_memory/unregister_host_memory. + return 4096 + + def register_host_memory(self, address, num_bytes): + # Register natively pinned (posix_memalign + mlock) host memory with the + # ACL runtime so torch's async copies can use the DMA engine. npurt + # initializes the runtime itself, so no set_device ordering is needed. + funcs, reason = _npu_host_copy_funcs() + if funcs is None: + from deepspeed.utils import logger + logger.warning_once(f"Host-memory registration is unavailable ({reason}); " + "native pinned memory stays mlock-only.") + return False + register, _ = funcs + rc = register(address, num_bytes, ACL_HOST_REG_MAPPED) + if rc != ACL_SUCCESS: + from deepspeed.utils import logger + logger.warning_once(f"npuHostRegister failed with rc={rc}; native pinned memory stays mlock-only.") + return False + return True + + def unregister_host_memory(self, address): + funcs, _ = _npu_host_copy_funcs() + if funcs is None: + return None + _, unregister = funcs + rc = unregister(address) + if rc != ACL_SUCCESS: + # Raise so NativePinnedMemory keeps the allocation alive: the driver + # must never hold a registration for pages later reused by malloc. + raise RuntimeError(f"npuHostUnregister failed with rc={rc}") + # Data types def is_bf16_supported(self): return torch.npu.is_bf16_supported() diff --git a/benchmarks/pin_memory/h2d_d2h_bench.py b/benchmarks/pin_memory/h2d_d2h_bench.py index 0cad917cdbb9..8211e351ea41 100644 --- a/benchmarks/pin_memory/h2d_d2h_bench.py +++ b/benchmarks/pin_memory/h2d_d2h_bench.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team -"""Compare torch and native pinned-memory H2D/D2H bandwidth on one CUDA GPU.""" +"""Compare torch and native pinned-memory H2D/D2H bandwidth on one accelerator device.""" import argparse import json @@ -55,7 +55,9 @@ def _time_copy(accelerator, copy_fn, stream, warmup, iters): def _allocate_host(accelerator, numel, arm): if arm == "torch": - return torch.empty(numel, dtype=torch.float32, pin_memory=True) + # Go through the accelerator so every backend pins against its own + # device rather than relying on torch's ``pin_memory=True`` fast path. + return accelerator._torch_pin_memory(torch.empty(numel, dtype=torch.float32)) return accelerator.pin_memory(torch.empty(numel, dtype=torch.float32), make_copy=False) @@ -65,8 +67,8 @@ def _run_arm(args): os.environ[key] = value accelerator = get_accelerator() - if accelerator.device_name() != "cuda" or not accelerator.is_available(): - raise RuntimeError("CUDA GPU is required") + if not accelerator.is_available(): + raise RuntimeError(f"No {accelerator.device_name()} device is available") accelerator.set_device(0) stream = accelerator.Stream() @@ -86,7 +88,7 @@ def _run_arm(args): "size_mib": size_mib, "h2d_gbps": num_bytes / h2d_seconds / 1e9, "d2h_gbps": num_bytes / d2h_seconds / 1e9, - "torch_is_pinned": host.is_pinned(), + "torch_is_pinned": accelerator._torch_is_pinned(host), "accelerator_is_pinned": accelerator.is_pinned(host), } print(f"RESULT={json.dumps(result, sort_keys=True)}", flush=True) diff --git a/deepspeed/utils/pin_memory.py b/deepspeed/utils/pin_memory.py index fa992df5c965..5099f6d43587 100644 --- a/deepspeed/utils/pin_memory.py +++ b/deepspeed/utils/pin_memory.py @@ -64,7 +64,13 @@ def _new_locked(self, example, out_shape): if base.nbytes and self._device_registration_enabled(): from deepspeed.accelerator import get_accelerator try: - if get_accelerator().register_host_memory(begin, base.nbytes): + accelerator = get_accelerator() + # Some device runtimes only register page-aligned ranges; round + # the address down to the accelerator-declared alignment and pad + # the size so the registered range still covers the request. + registered_begin = self._align_host_address(accelerator, begin) + padded_bytes = base.nbytes + (begin - registered_begin) + if accelerator.register_host_memory(registered_begin, padded_bytes): self._device_registered.add(begin) except Exception as e: logger.warning_once( @@ -131,12 +137,24 @@ def _release(handle, begin, ranges, finalizers, device_registered): # during interpreter shutdown. pass + @staticmethod + def _align_host_address(accelerator, address): + # Device runtimes may require page-aligned registration ranges; round + # down to the alignment the accelerator declares (1 = no requirement). + alignment = accelerator.pin_memory_alignment() + if alignment <= 1: + return address + return address - (address % alignment) + @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) + accelerator = get_accelerator() + # Same rounding as at registration time, so the driver releases exactly + # the range it was given. + accelerator.unregister_host_memory(NativePinnedMemory._align_host_address(accelerator, begin)) device_registered.discard(begin) @staticmethod diff --git a/docs/code-docs/source/memory.rst b/docs/code-docs/source/memory.rst index 1a52b04cc3f2..614824e0c59a 100644 --- a/docs/code-docs/source/memory.rst +++ b/docs/code-docs/source/memory.rst @@ -411,10 +411,13 @@ Example: 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: +Native allocations are device-independent ``mlock`` buffers. DeepSpeed +additionally registers them with the device through the accelerator's +host-memory registration hook, so PyTorch can use them for asynchronous +H2D/D2H DMA. Accelerators declare the alignment required for registered host +memory via ``pin_memory_alignment``; native buffers are rounded down and +size-padded to it before registration. Device registration is enabled by +default and can be disabled for comparison or debugging: .. code-block:: bash diff --git a/tests/unit/v1/pin_memory/test_pin_memory.py b/tests/unit/v1/pin_memory/test_pin_memory.py index d30f38a37706..20895b0ed554 100644 --- a/tests/unit/v1/pin_memory/test_pin_memory.py +++ b/tests/unit/v1/pin_memory/test_pin_memory.py @@ -5,8 +5,10 @@ import pytest import torch +from deepspeed.accelerator import npu_accelerator from deepspeed.accelerator.cpu_accelerator import CPU_Accelerator from deepspeed.accelerator.cuda_accelerator import CUDA_Accelerator +from deepspeed.accelerator.npu_accelerator import NPU_Accelerator from deepspeed.utils.pin_memory import NativePinnedMemory @@ -140,6 +142,9 @@ def register_host_memory(self, address, num_bytes): def unregister_host_memory(self, address): self.unregistered.append(address) + def pin_memory_alignment(self): + return 1 + def test_native_device_registration_and_unpin(monkeypatch, native_pins): accelerator = _RegisteringAccelerator() @@ -225,6 +230,9 @@ def register_host_memory(self, address, num_bytes): def unregister_host_memory(self, address): raise AssertionError("unregister must not run when register failed") + def pin_memory_alignment(self): + return 1 + 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) @@ -284,3 +292,157 @@ def unregister_host_memory(self, address): assert native_pins.unpin(pinned) is True assert accelerator.unregistered == [begin] assert begin not in native_pins._device_registered + + +class _AlignedAccelerator(_RegisteringAccelerator): + + def __init__(self, alignment): + super().__init__() + self._alignment = alignment + + def pin_memory_alignment(self): + return self._alignment + + +class _OffsetHandle: + """pin_memory-op stand-in whose buffer base sits at a fixed byte offset.""" + + def __init__(self, offset): + self._offset = offset + + def new_cpu_locked_tensor(self, numel, example): + storage = torch.empty(numel * example.element_size() + self._offset, dtype=torch.uint8) + return storage[self._offset:].view(example.dtype) + + def free_cpu_locked_tensor_by_ptr(self, address): + return True + + +@pytest.mark.parametrize("alignment, offset", [(1, 64), (2048, 1232), (4096, 64), (4096, 0)]) +def test_device_registration_aligns_to_declared_alignment(native_pins, monkeypatch, alignment, offset): + # Accelerators declare the alignment their device runtime requires for + # host-memory registration. NativePinnedMemory must round the registered + # range down to it, pad the size so the full request is covered, and + # unregister the same aligned address. Alignment 1 means no requirement, + # so the request passes through unchanged. + accelerator = _AlignedAccelerator(alignment) + monkeypatch.setattr("deepspeed.accelerator.get_accelerator", lambda: accelerator) + monkeypatch.setenv("DS_PIN_MEMORY_REGISTER_DEVICE", "1") + monkeypatch.setattr(native_pins, "_handle", _OffsetHandle(offset)) + + pinned = native_pins.pin(torch.empty(32), make_copy=False) + begin = pinned.data_ptr() + registered_address, registered_bytes = accelerator.registered[0] + if alignment == 1: + assert registered_address == begin + else: + assert registered_address % alignment == 0 + assert registered_address <= begin < registered_address + alignment + assert registered_bytes == pinned.nbytes + (begin - registered_address) + + assert native_pins.unpin(pinned) is True + assert accelerator.unregistered == [registered_address] + + +def test_npu_declares_page_alignment(): + # MAPPED registration rejects non-4K-aligned addresses; the declared + # alignment is what makes NativePinnedMemory round ranges down for NPU. + accelerator = NPU_Accelerator.__new__(NPU_Accelerator) + assert accelerator.pin_memory_alignment() == 4096 + + +def test_npu_register_uses_mapped_flag(monkeypatch): + # MAPPED is deliberate: PINNED-only registrations fall back to mlock-speed + # copies on this platform (measured ~9 GB/s vs ~23 GB/s for 64 MiB buffers). + registered = [] + unregistered = [] + + def register(addr, num_bytes, flag): + registered.append((addr, num_bytes, flag)) + return 0 + + def unregister(addr): + unregistered.append(addr) + return 0 + + monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None)) + accelerator = NPU_Accelerator.__new__(NPU_Accelerator) + + assert accelerator.register_host_memory(4096, 4096) is True + assert registered == [(4096, 4096, npu_accelerator.ACL_HOST_REG_MAPPED)] + accelerator.unregister_host_memory(4096) + assert unregistered == [4096] + + +def test_npu_device_registration_failure_returns_false(monkeypatch): + # A non-zero npuHostRegister return code must degrade to mlock-only, not + # raise: the NativePinnedMemory caller only tracks the address on True. + def register(address, num_bytes, flag): + return 107000 + + def unregister(address): + raise AssertionError("unregister must not run when register failed") + + monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None)) + accelerator = NPU_Accelerator.__new__(NPU_Accelerator) + + assert accelerator.register_host_memory(4096, 4096) is False + + +def test_npu_unregister_failure_raises(monkeypatch): + # Raising keeps the allocation alive in NativePinnedMemory so the driver + # never holds a registration for pages later reused by malloc. + def register(address, num_bytes, flag): + return 0 + + def unregister(address): + return 107000 + + monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: ((register, unregister), None)) + accelerator = NPU_Accelerator.__new__(NPU_Accelerator) + + with pytest.raises(RuntimeError, match="npuHostUnregister"): + accelerator.unregister_host_memory(4096) + + +def test_npu_missing_npurt_is_noop(monkeypatch): + monkeypatch.setattr(npu_accelerator, "_npu_host_copy_funcs", lambda: (None, "test")) + accelerator = NPU_Accelerator.__new__(NPU_Accelerator) + + assert accelerator.register_host_memory(4096, 4096) is False + assert accelerator.unregister_host_memory(4096) is None + + +@pytest.mark.skipif(not hasattr(torch, "npu") or not hasattr(torch.npu, "npurt"), reason="torch_npu is not installed") +def test_npu_host_copy_lookup_gates(monkeypatch): + # Exercise each npurt resolution path (missing, init failure, success) + # by stubbing torch.npu; where torch_npu is absent the whole test is + # skipped because stubbing torch.npu is not reliable across torch versions. + class _StubNpu: + pass + + monkeypatch.setattr(torch, "npu", _StubNpu(), raising=False) + # A build lacking npurt must not resolve, with a reason saying so. + monkeypatch.delattr(torch.npu, "npurt", raising=False) + funcs, reason = npu_accelerator._npu_host_copy_funcs() + assert funcs is None + assert "npurt is unavailable" in reason + + # npurt() failing to initialize the runtime resolves to None with a reason. + def fail_npurt(): + raise RuntimeError("init failed") + + monkeypatch.setattr(torch.npu, "npurt", fail_npurt, raising=False) + funcs, reason = npu_accelerator._npu_host_copy_funcs() + assert funcs is None + assert "initialize" in reason + + # A working npurt module resolves to its host copy functions. + class _Npurt: + npuHostRegister = "register" + npuHostUnregister = "unregister" + + monkeypatch.setattr(torch.npu, "npurt", lambda: _Npurt(), raising=False) + funcs, reason = npu_accelerator._npu_host_copy_funcs() + assert funcs == ("register", "unregister") + assert reason is None