[REQUEST] Register native pinned host memory with device runtime for NPU (follow-up to #8283 and #8315)
Is your feature request related to a problem? Please describe.
DeepSpeed's native pinned-memory backend (DS_PIN_MEMORY_BACKEND=native) page-locks
host buffers with posix_memalign + mlock, and since #8283 it can additionally
register each buffer with the device runtime so asynchronous copies can use the DMA
engine. That change made a large difference where it shipped: on an H200, H2D
bandwidth went from ~10-17 GB/s (mlock alone) to ~53-55 GB/s after registration.
The registration capability, however, is a per-accelerator hook, and today only the
CUDA accelerator implements it. Every other platform falls through to an empty
default: with DS_PIN_MEMORY_REGISTER_DEVICE=1 on an Ascend NPU, the hook returns
False silently, the buffer remains mlock-only, and the native backend performs
like pageable memory. This is not a hypothetical gap — on an Ascend 910B4
(CANN 8.5.0) registering the same native buffer lifts 4 MiB H2D from ~10 GB/s to
~23 GB/s, and 64 MiB registered copies hold ~23.4 GB/s while mlock-only
fluctuates between ~13 and ~22.5 GB/s.
The XPU platform faces the identical issue, already filed as #8315. That NPU and XPU users are independently hitting the same blind spot makes the point: #8283's benefit is gated per platform, and the non-CUDA platforms that run an increasing share of real training jobs (Huawei Ascend in particular) silently lose 2x+ on every H2D copy until someone implements the two hooks for them. This issue covers the NPU side; #8315 covers XPU.
Describe the solution you'd like
Implement register_host_memory / unregister_host_memory on NPU_Accelerator
(accelerator/npu_accelerator.py) via the NPU runtime binding torch_npu ships:
torch.npu.npurt()returns the NPU runtime-API module withnpuHostRegister(addr, size, flags)/npuHostUnregister(addr), thin wrappers overaclrtHostRegisterV2/aclrtHostUnregister. It initializes the runtime if needed and fails with an explicitRuntimeErrorwhen the build does not include it.- Failures should not disturb training: if torch_npu is older than 2.9.0, or the
build has no npurt, or the registration call returns a non-zero code, log a
warning once and continue with plain mlock memory (the existing
DS_PIN_MEMORY_REGISTER_DEVICEswitch keeps acting as the opt-out). - Extend
tests/unit/v1/pin_memory/test_pin_memory.pywith NPU cases mirroring the CUDA cudart-mock test, plus coverage for the availability fallbacks. - The before/after bandwidth can be measured on NPU with
benchmarks/pin_memory/h2d_d2h_bench.pyonce it is made accelerator-agnostic.
Describe alternatives you've considered
- Leave NPU as-is: native pin stays mlock-only on Ascend, which is functionally correct but pays the full pageable-copy penalty this feature was created to remove.
- Give the NPU native backend an
aclrtMallocHost-style allocator: true pinned allocation instead of post-hoc registration, but it would replace the device-independent allocator wholesale and is a much bigger change than the two-hook contract #8283 established.
Additional context
- Hook definitions and call sites: #8283 introduced the hooks (empty defaults in
accelerator/abstract_accelerator.py, CUDA implementation viacudaHostRegister/cudaHostUnregister) and invokes them fromNativePinnedMemory.pin/_releaseindeepspeed/utils/pin_memory.py. - XPU counterpart: #8315, with an implementation attempt in PR #8322 (blocked on oneAPI 2026.2). Resolving both follow-ups completes the picture for the non-CUDA accelerators DeepSpeed supports.
- Validation environment: CANN 8.5.0, torch 2.10.0+cpu, torch_npu 2.10.0, Ascend 910B4 (aarch64).
Source: deepspeedai/DeepSpeed