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

[sfc-gh-truwase]

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 DeepNVMe async-io (mlock) allocator instead of torch.pin_memory().

Depends on / builds atop #8212 (shared DeepNVMe pinned-tensor manager). Native allocations go through new_cpu_locked_tensor, so I/O handles recognize them via the process-wide manager and skip bounce buffers.

• 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 async-io 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.
• Consolidation: XPU's bespoke native-pinning path is folded into the shared backend.
• Callers: compile paths route through get_accelerator(). Swap-tensor buffers continue to allocate via I/O handles (from Share DeepNVMe pinned-tensor manager and route swap buffers through I/O handles #8212); 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); also pushed to rtd-staging.
• Tests: unit tests for the native manager, accelerator pinning APIs, and destroy-path unpin; accelerator tests moved under tests/unit/v1.

Test plan

• pre-commit (yapf, flake8, check-license, check-torchdist, codespell) passes on all changed files.
• Rebased onto master after Share DeepNVMe pinned-tensor manager and route swap buffers through I/O handles #8212 merge; shared-manager overlap resolved by keeping Share DeepNVMe pinned-tensor manager and route swap buffers through I/O handles #8212's swap/I/O paths.
• Focused pin UTs on a GPU node: tests/unit/v1/pin_memory/test_pin_memory.py + tests/unit/v1/accelerator/test_accelerator.py + tests/unit/v1/nvme/test_pinned_manager.py (incl. GC finalizer coverage) — 21 passed.
• Destroy-path unpin UTs: tests/unit/v1/pin_memory/test_destroy_unpin.py — 5/5 passed (native ZeRO-2/3 free all optimizer-owned ranges; native ZeRO-3+param-offload decreases ranges; torch backend destroy is a clean no-op).
• Full tests/unit/v1/ suite on a GPU node — 1134 passed, 63 skipped, exit 0 (~2.5h).
• Bounce-buffer check: under DS_PIN_MEMORY_BACKEND=native, a buffer from get_accelerator().pin_memory (and a narrow of it) is is_pinned on a different AIO handle than the allocator.
• RTD staging render of Host Memory Pinning: https://deepspeed.readthedocs.io/en/rtd-staging/memory.html#host-memory-pinning
• CPU-node sanity check on tunji-cpu-ds-0 (DS_ACCELERATOR=cpu): torch backend 11 passed; native backend 11 passed (test_pin_memory.py, native accelerator APIs, test_pinned_manager.py). XPU/HPU still pending reviewer hardware.

Made with Cursor

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: cee4d7771e

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[delock]

I think host memory pinning accouting is needed. I made an attempt in the following PR (in order to justfy pin_memory as default offload option).
https://github.com/deepspeedai/DeepSpeed/pull/8207/changes

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f7934bc8d3

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[sfc-gh-truwase]

Addressed Codex review (r3713658448) by landing #8212 first, then rebasing this PR onto updated master (f7934bc8).

With the process-wide shared deepspeed_pin_tensor_t manager, buffers allocated via NativePinnedMemory (new_cpu_locked_tensor) are registered in the same manager every I/O handle consults. Verified on a GPU node: a native-pinned buffer (and a narrow of it) is reported is_pinned by a different AIO handle than the allocator, so DeepNVMe will skip the bounce-buffer path.

[sfc-gh-truwase]

Addressed the native pin lifetime concern with two complementary paths:

1. weakref.finalize on NativePinnedMemory.pin() (already landed) so GC frees mlocked allocations when tensors are dropped without an explicit unpin.
2. Explicit accelerator.unpin_memory() in ZeRO/ZenFlow destroy() for optimizer-owned CPU-offload pins (this commit), so teardown does not wait on GC.

Coverage:

• ZeRO-1/2, ZeRO-3, and ZenFlow overrides unpin their owned offload buffers in destroy().
• Param-offload partition pins that are rebound to flattened views are left to the finalizer (unpinning those views is ineffective / unsafe).
• New UTs in tests/unit/v1/pin_memory/test_destroy_unpin.py (5/5 passed on GPU): native ZeRO-2/3 free all optimizer pins; native ZeRO-3+param-offload decreases ranges; torch backend destroy remains a clean no-op.

[sfc-gh-truwase]

@delock , @stas00, @tohtana see this page for docs preview of this feature:
https://deepspeed.readthedocs.io/en/rtd-staging/memory.html#host-memory-pinning

[sfc-gh-truwase]

@delock Accounting from #8207 is now integrated into this PR (merged latest master in 0cca5393).

get_accelerator().pin_memory() always goes through track_pinned_memory for both the default torch backend and DS_PIN_MEMORY_BACKEND=native (native tracks the returned locked buffer’s nbytes). The CPU accelerator still bypasses accounting on its torch no-op path (nothing is page-locked), and defers to the ABC — with tracking — when the native backend is active.

Pinned-memory volume notes in the RTD Memory Usage page also include the #8207 updates, alongside the Host Memory Pinning section for backend selection.