DeepSpeed

aa0e91b9 - Add native (DeepNVMe) host-memory pinning backend for accelerators (#8211)

Add native (DeepNVMe) host-memory pinning backend for accelerators (#8211) ## Summary Adds a native host-memory pinning backend, selectable via the `DS_PIN_MEMORY_BACKEND` environment variable (defaults to `torch`). When set to `native`, CPU memory is page-locked through the standalone DeepSpeed `pin_memory` op (`PinMemoryBuilder` / `pin_handle`, `posix_memalign` + `mlock`) instead of `torch.pin_memory()`. Stacked on #8236 (standalone `pin_memory` op, now on `master`). Native allocations go through `pin_handle`, so DeepNVMe I/O handles recognize them via the process-wide manager and skip bounce buffers — without requiring libaio / AIO worker threads. - **New `deepspeed/utils/pin_memory.py`**: a process-wide shared `NativePinnedMemory` manager that pins CPU memory, tracks pinned pointer ranges (so slices/views report as pinned), tags buffers with `.ds_pinned`, supports `make_copy`/`match_shape`, and frees on unpin. It fails early with a clear error if the `pin_memory` op cannot be built (no silent torch fallback). Native pins also use a `weakref` finalizer so GC releases mlocked pages when tensors are dropped without an explicit unpin. - **Accelerator owns dispatch**: `pin_memory` drops `align_bytes` and gains `make_copy`/`match_shape`; `is_pinned` is FakeTensor/meta-tensor safe; new `unpin_memory` (native frees, torch no-op). Subclasses retain only the device-specific `_torch_pin_memory`/`_torch_is_pinned` primitives. Preserves master's `track_pinned_memory` accounting (CPU torch no-op still bypasses it). - **Consolidation**: XPU's bespoke `align_bytes=0` path is folded into the shared native backend. - **Callers**: `compile` paths route through `get_accelerator()`. Swap-tensor buffers continue to allocate via I/O handles; with the shared manager they interoperate with native-pinned tensors. ZeRO / ZenFlow `destroy()` explicitly unpins optimizer-owned CPU-offload buffers under the native backend. - **Docs**: Host Memory Pinning section under RTD Memory Usage (`docs/code-docs/source/memory.rst`). - **Tests**: unit tests for the native manager, accelerator pinning APIs, destroy-path unpin, and cross-op recognition with AIO. ## Test plan - [x] Rebased onto `master` after #8236 merge; retargeted `NativePinnedMemory` from `AsyncIOBuilder` → `PinMemoryBuilder`. - [x] `pre-commit` on changed files. - [x] Focused UTs on GPU (`tunji-h200-n1g2-ds2-0`, job `20260809T183311Z`): `tests/unit/v1/pin_memory/` + `tests/unit/v1/accelerator/test_accelerator.py` + `tests/unit/v1/nvme/test_pinned_manager.py` — **30 passed**. - [x] Bounce-buffer / cross-op smoke: under `DS_PIN_MEMORY_BACKEND=native`, a `pin_handle` buffer is `is_pinned` on a separate AIO handle. Made with [Cursor](https://cursor.com) --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>