#8531·DeepSpeed

[REQUEST] Register native pinned host memory with device runtime for NPU (follow-up to #8283 and #8315)

Author: VenusTZZCreated Sep 16, 2026Updated Sep 16, 2026
Labelsenhancement

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 with npuHostRegister(addr, size, flags) / npuHostUnregister(addr), thin wrappers over aclrtHostRegisterV2 / aclrtHostUnregister. It initializes the runtime if needed and fails with an explicit RuntimeError when 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_DEVICE switch keeps acting as the opt-out).
  • Extend tests/unit/v1/pin_memory/test_pin_memory.py with 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.py once 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 via cudaHostRegister / cudaHostUnregister) and invokes them from NativePinnedMemory.pin / _release in deepspeed/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).