DeepSpeed

0ca14e83 - Register native pinned host memory with the CANN runtime for NPU (#8532)

Register native pinned host memory with the CANN runtime for NPU (#8532) Fixes #8531. Follow-up to #8283; XPU counterpart tracked in #8315 (attempted in #8322). ## Summary - Implement `register_host_memory` / `unregister_host_memory` on `NPU_Accelerator`via the NPU runtime API torch_npu ships: `torch.npu.npurt()` returns the runtime-API module with `npuHostRegister` / `npuHostUnregister`, thin wrappers over `aclrtHostRegisterV2` / `aclrtHostUnregister`. Previously `NPU_Accelerator` inherited the base-class no-ops, so `DS_PIN_MEMORY_BACKEND=native` buffers stayed mlock-only on Ascend NPUs and asynchronous copies ran at pageable bandwidth. - Degrades to "log once and continue with mlock" on every failure path and keeps honoring the existing `DS_PIN_MEMORY_REGISTER_DEVICE` opt-out. - Extends `tests/unit/v1/pin_memory/test_pin_memory.py` with NPU cases mirroring the CUDA cudart-mock tests, including the availability fallbacks. - Makes `benchmarks/pin_memory/h2d_d2h_bench.py` accelerator-agnostic (it hardcoded a CUDA check and torch's `pin_memory=True` fast path) so the before/after bandwidth could be measured on NPU. - Documents the NPU registration path in `docs/code-docs/source/memory.rst`. ## Implementation notes **The npurt binding.** torch_npu exposes `torch.npu.npurt()`, which initializes the runtime if needed and returns the runtime-API module. Using it avoids binding `libascendcl.so` ourselves and raises a clear `RuntimeError` when the build lacks the module. Two defenses gate the hook, both falling back to mlock-only with a one-time warning: - no `torch.npu.npurt` attribute (the hook checks for its presence), - `npurt()` failing to initialize the runtime. **Flag choice: `ACL_HOST_REG_MAPPED`.** The `flags` parameter of `npuHostRegister` is the CANN `ACL_HOST_REG_*` bitmask (per the npurt docstring and `acl_rt.h`): `MAPPED` (0x2) page-locks the range and adds a device mapping, `PINNED` (0x10000000) only page-locks. Measured through the bench methodology at 64 MiB (10 process-level A/B runs): `PINNED`-only registrations fell back to mlock-only H2D speed (~8-9 GB/s) in 5/5 runs, while `MAPPED` kept ~23 GB/s (5/6). `MAPPED` additionally requires 4K-aligned addresses (a 64B-offset probe fails with `ACL_ERROR_RT_DRV_INTERNAL_ERROR`, 507899), which DeepSpeed's native allocator guarantees (`posix_memalign(..., sysconf(_SC_PAGESIZE), ...)`); sizes need no page alignment (probed 400 B to 1 MiB+3, rc=0). As a guard against unaligned callers, the hook extends the range down to the page boundary, padding the size accordingly, instead of relying on the driver's opaque internal error. **Failure semantics.** A non-zero return code from `npuHostRegister` logs once and returns `False`, so `NativePinnedMemory` keeps the buffer mlock-only and does not track the address. A failed `npuHostUnregister` raises (like the CUDA hook), which makes `NativePinnedMemory` keep the allocation alive so the driver never holds a registration for pages later reused by `malloc`. ## Bandwidth on Ascend 910B4 `benchmarks/pin_memory/h2d_d2h_bench.py`, one full run of the #8283 methodology (`--sizes-mib 4 64 --warmup 10 --iters 50`), CANN 8.5.0, torch 2.10.0+cpu, torch_npu 2.10.0, Ascend 910B4 (aarch64), shared box: | Arm | Size (MiB) | H2D (GB/s) | D2H (GB/s) | torch.is_pinned | |-----|------------|------------|------------|-----------------| | torch | 4 | 10.4 | 21.1 | True | | torch | 64 | 16.5 | 21.0 | True | | native-unregistered | 4 | 9.7 | 9.4 | False | | native-unregistered | 64 | 21.1 | 19.6 | False | | native-registered | 4 | **22.7** | 20.6 | True | | native-registered | 64 | **23.4** | 20.9 | True | After registration, `torch.is_pinned()` on a `torch.frombuffer` view of the range returns `True` — the torch_npu runtime recognizes the registered range — so the DMA path is confirmed to engage, not just inferred from the numbers. Registration lifts H2D bandwidth from ~9.7 to ~22.7 GB/s at 4 MiB (2.3x) and holds ~23.4 GB/s at 64 MiB, with D2H improving as well (9.4 → 20.6 and 19.6 → 20.9 GB/s). Registered copies match or exceed the torch `pin_memory()` arm at every size. ## Test plan - [x] `pre-commit run --files` on touched paths — yapf, flake8, check-torchdist, check-torchcuda, check-license, codespell all pass - [x] `pytest tests/unit/v1/pin_memory/test_pin_memory.py` — **24 passed** on Ascend 910B4 (18 existing incl. the CUDA cudart-mock tests + 6 new NPU cases: npurt register/unregister calls, unaligned address extended to the page boundary, non-zero rc degrades to mlock, unregister failure raises, missing npurt is a no-op, availability gate table) - [x] `benchmarks/pin_memory/h2d_d2h_bench.py` on Ascend 910B4 — all three arms run to completion and exit 0 at 4/64 MiB ## Known limitation Some torch_npu builds may ship without the npurt module. The `hasattr` check degrades to mlock-only with a one-time warning, so registration is only available when the binding is present. --------- Signed-off-by: Zizhao Tang <1601677496@qq.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>