Add native (DeepNVMe) host-memory pinning backend for accelerators - #8211
Add native (DeepNVMe) host-memory pinning backend for accelerators#8211sfc-gh-truwase wants to merge 7 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
|
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). |
cee4d77 to
f7934bc
Compare
There was a problem hiding this comment.
💡 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".
|
Addressed Codex review (r3713658448) by landing #8212 first, then rebasing this PR onto updated With the process-wide shared |
|
Addressed the native pin lifetime concern with two complementary paths:
Coverage:
|
|
@delock , @stas00, @tohtana see this page for docs preview of this feature: |
|
@delock Accounting from #8207 is now integrated into this PR (merged latest
Pinned-memory volume notes in the RTD Memory Usage page also include the #8207 updates, alongside the Host Memory Pinning section for backend selection. |
|
Converted to draft. This PR will be rebased/stacked on top of #8236 (standalone |
Host Memory Pinning now covers only torch/accelerator pinning and the standalone pin_memory op; DS_PIN_MEMORY_BACKEND=native belongs in #8211. Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
…8236) ## Summary - Extract host page-locking (`posix_memalign`/`mlock`) into a standalone `deepspeed.ops.pin_memory` / `PinMemoryBuilder` that does not require libaio or AIO worker threads. - Compile the pin manager only in the `pin_memory` op and share one process-wide manager with `async_io`/`gds` (via exported symbol + `RTLD_GLOBAL`) so DeepNVMe bounce-buffer skipping and `is_pinned` stay consistent. - Keep `new_cpu_locked_tensor` / `free_cpu_locked_tensor` / `is_pinned` on `aio_handle`/`gds_handle` as thin wrappers; point XPU `align_bytes=0` at `pin_handle`. ## Test plan - [x] `tests/unit/v1/pin_memory/test_pin_memory_op.py` (pin without async_io) - [x] `tests/unit/v1/nvme/test_pinned_manager.py` cross-op recognition (`pin_handle` ↔ `aio_handle`) - [x] Confirm AIO I/O tests still pass with shared manager (`tests/unit/v1/nvme/` including `test_aio.py` / `test_gds.py` — 131 passed on `tunji-h200-n1g2-ds2-0`, job `20260809T123310Z`, HEAD `a6f6ab6e`) - [x] `ds_report` shows `pin_memory` as compatible without libaio-dev (`pin_memory ... [OKAY]`; with `io_submit`/libaio mocked missing, `pin_memory` stays compatible while `async_io` does not — job `20260809T124859Z`) ## Follow-up Native backend (`DS_PIN_MEMORY_BACKEND=native`, deepspeedai#8211) will be stacked on this PR once it lands. Made with [Cursor](https://cursor.com) --------- Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
Introduce a native host-memory pinning backend selectable via the DS_PIN_MEMORY_BACKEND env var (defaults to "torch"). When set to "native", CPU memory is page-locked through the DeepNVMe async-io (mlock) allocator instead of torch.pin_memory(). - deepspeed/utils/pin_memory.py: new NativePinnedMemory manager (process-wide shared singleton) 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. Construction fails early with a clear error if the async-io op cannot be built. - Accelerator owns the native-vs-torch dispatch. pin_memory drops align_bytes and gains make_copy/match_shape; is_pinned is FakeTensor/meta-tensor safe; add unpin_memory (native frees, torch is a no-op). Subclasses keep only the device-specific _torch_pin_memory/_torch_is_pinned primitives. - Consolidate XPU's bespoke native pinning path into the shared backend. - Route swap_tensor and compile callers through the accelerator. - Add unit tests for the native manager and the accelerator pinning APIs, and move the accelerator tests under tests/unit/v1. Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
Native pin_memory() allocations are page-locked through the AIO manager and were only released by an explicit unpin_memory(), which existing callers relying on torch.pin_memory()'s GC-based lifetime never call -- so repeatedly creating and dropping pinned buffers leaked mlocked host memory until process exit. Attach a weakref.finalize to the returned tensor that frees the allocation when it is garbage-collected, and cancel that finalizer on explicit unpin() to avoid a redundant free. Addresses Codex review r3721965881. Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
Release optimizer-owned page-locked host buffers when engines tear down under DS_PIN_MEMORY_BACKEND=native, so mlocked memory does not wait on GC. Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
Document DS_PIN_MEMORY_BACKEND (torch vs native), the async-io build requirement, and native unpin/lifetime semantics for the new pinning backend. Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
Colocate NativePinnedMemory unit tests with the destroy-path pin_memory suite under the v1 tree. Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
NativePinnedMemory now uses PinMemoryBuilder/pin_handle instead of AsyncIOBuilder, matching the #8236 split so native mode no longer requires libaio. Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
c1f11be to
1bdd3d2
Compare
|
Resumed after #8236 landed:
Ready for review. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1bdd3d2d45
ℹ️ 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".
| self._finalizers[begin] = weakref.finalize(locked, self._release, self._handle, base, self._ranges, | ||
| self._finalizers, begin) |
There was a problem hiding this comment.
With DS_PIN_MEMORY_BACKEND=native, if a caller keeps a slice/view of the returned tensor but drops the returned tensor, this finalizer is attached only to locked, so it can run while the derived view still aliases the same storage. The fresh issue is that the patch explicitly supports views as pinned, but the callback frees base, leaving that live view usable by ZeRO/AIO after the mlocked allocation has been freed; tie cleanup to the shared storage/all aliases rather than just the returned view.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7df48c5. The weakref.finalize is now tied to the allocation root (base), not the returned view. Every slice/view keeps base alive through ._base, so the mlocked region is freed only after the returned tensor and all its aliases are gone — a live view can no longer alias a freed allocation. Freeing is done by address (see below) so base is not pinned alive by the finalizer itself. Added test_allocation_survives_until_all_views_dropped.
|
|
||
| def unpin(self, tensor): | ||
| # After freeing, using ``tensor`` is a use-after-free and must be avoided. | ||
| begin = tensor.data_ptr() |
There was a problem hiding this comment.
When a native-pinned tensor has its .data redirected before destroy() (for example Stage-3 offload_states() assigns grad_partitions_flat_buffer.data to an LP-grad pin buffer and reload_states() can move it back to device), using the current data_ptr() loses the original allocation key and finalizer. In that context destroy() frees the current/no buffer and leaves the original mlocked range in _ranges until the whole tensor is garbage-collected, so engines that offload/reload states do not release native pins deterministically.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7df48c5. pin() now records the owning allocation address on the returned tensor (ds_pin_base), and both unpin() and the finalizer free by that address via the new pin_handle.free_cpu_locked_tensor_by_ptr / deepspeed_pin_tensor_t::free(void*). So even after Stage-3 offload_states() redirects grad_partitions_flat_buffer.data, destroy() releases the original locked range deterministically instead of leaving it until GC. Added test_unpin_frees_original_after_data_redirect.
| @@ -0,0 +1,122 @@ | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
There was a problem hiding this comment.
This is a non-merge commit (git rev-list --parents -n 1 shows a single parent), but the commit message has no Signed-off-by: trailer. The repository policy requires every non-merge commit to include one, so DCO/policy checks can block the PR until the commit is recreated with --signoff.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
All commits on this branch carry a Signed-off-by trailer (created with --signoff); the DCO / required check is green. This looks like a false positive from reviewing the diff in isolation.
| def LongTensor(self): | ||
| return functools.partial(torch.tensor, dtype=torch.long, device=self._name) | ||
|
|
||
| def _pin_memory(self, tensor, align_bytes=1): |
There was a problem hiding this comment.
Is this mean XPU accelerator needs to explicitly set DS_PIN_MEMORY_BACKEND to native with intended page-locked host buffer?
@rogerxfeng8 to confirm wheter XPU still relies on new_cpu_locked_tensor to work in certain scenario.
There was a problem hiding this comment.
@delock @rogerxfeng8 yes it will be great to confirm the best support for XPU.
Address Codex review on the native pinning backend: - Tie the GC finalizer to the allocation root rather than the returned view, so a live slice/view can no longer outlive the freed mlocked region (previously dropping the returned tensor while a view was alive caused a use-after-free). Views keep the root alive via ._base, so the allocation is released only after every alias is gone. - Record the owning allocation address (ds_pin_base) at pin time and free by address in both unpin() and the finalizer, so an explicit unpin() releases the original range even after a tensor's .data is redirected (e.g. ZeRO offload/reload rebinds .data to a different buffer). Add pin_handle.free_cpu_locked_tensor_by_ptr / deepspeed_pin_tensor_t::free(void*) to support address-based release, and cover both cases with unit tests. Signed-off-by: Olatunji Ruwase <[email protected]> Co-authored-by: Cursor <[email protected]>
2 participants
Summary
Adds a native host-memory pinning backend, selectable via the
DS_PIN_MEMORY_BACKENDenvironment variable (defaults totorch). When set tonative, CPU memory is page-locked through the standalone DeepSpeedpin_memoryop (PinMemoryBuilder/pin_handle,posix_memalign+mlock) instead oftorch.pin_memory().Stacked on #8236 (standalone
pin_memoryop, now onmaster). Native allocations go throughpin_handle, so DeepNVMe I/O handles recognize them via the process-wide manager and skip bounce buffers — without requiring libaio / AIO worker threads.deepspeed/utils/pin_memory.py: a process-wide sharedNativePinnedMemorymanager that pins CPU memory, tracks pinned pointer ranges (so slices/views report as pinned), tags buffers with.ds_pinned, supportsmake_copy/match_shape, and frees on unpin. It fails early with a clear error if thepin_memoryop cannot be built (no silent torch fallback). Native pins also use aweakreffinalizer so GC releases mlocked pages when tensors are dropped without an explicit unpin.pin_memorydropsalign_bytesand gainsmake_copy/match_shape;is_pinnedis FakeTensor/meta-tensor safe; newunpin_memory(native frees, torch no-op). Subclasses retain only the device-specific_torch_pin_memory/_torch_is_pinnedprimitives. Preserves master'strack_pinned_memoryaccounting (CPU torch no-op still bypasses it).align_bytes=0path is folded into the shared native backend.compilepaths route throughget_accelerator(). Swap-tensor buffers continue to allocate via I/O handles; with the shared manager they interoperate with native-pinned tensors. ZeRO / ZenFlowdestroy()explicitly unpins optimizer-owned CPU-offload buffers under the native backend.docs/code-docs/source/memory.rst).Test plan
masterafter Split native host pinning into standalone pin_memory op #8236 merge; retargetedNativePinnedMemoryfromAsyncIOBuilder→PinMemoryBuilder.pre-commiton changed files.tunji-h200-n1g2-ds2-0, job20260809T183311Z):tests/unit/v1/pin_memory/+tests/unit/v1/accelerator/test_accelerator.py+tests/unit/v1/nvme/test_pinned_manager.py— 30 passed.DS_PIN_MEMORY_BACKEND=native, apin_handlebuffer isis_pinnedon a separate AIO handle.Made with Cursor