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".
| # surface the build/load failure here instead of silently degrading. | ||
| try: | ||
| from deepspeed.ops.op_builder import AsyncIOBuilder | ||
| self._handle = AsyncIOBuilder().load().aio_handle(128 * 1024, 8, False, False, 1) |
There was a problem hiding this comment.
When DS_PIN_MEMORY_BACKEND=native is used for non-GDS DeepNVMe swapping, this constructs a private AIO handle just for allocation, but the actual read/write handles each create their own _pinned_tensor_mgr and cpu_op_desc_t only skips the bounce buffer for tensors managed by that same handle (csrc/aio/py_lib/deepspeed_py_io_handle.cpp:49,253 and csrc/aio/py_lib/deepspeed_cpu_op.cpp:26-27). As a result, buffers returned by get_accelerator().pin_memory(..., make_copy=False) satisfy the Python range check but are not recognized by the handle performing I/O, so every native-pinned swap still gets copied through a bounce buffer; share the pinned manager or allocate through the I/O handle that will use the buffer.
Useful? React with 👍 / 👎.
|
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). |
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 <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
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".
|
|
||
| def pin(self, tensor, make_copy=True, match_shape=True): | ||
| numel = tensor.numel() | ||
| locked = self._handle.new_cpu_locked_tensor(numel, tensor)[:numel] |
There was a problem hiding this comment.
When DS_PIN_MEMORY_BACKEND=native is enabled, existing get_accelerator().pin_memory(...) call sites still rely on the old torch.pin_memory() lifetime semantics where dropping the tensor releases the pinned allocation. This allocation goes through the AIO pinned manager and is only freed via the new unpin_memory() path, which those existing call sites do not call, so repeatedly creating/destroying engines or offload buffers in one process can accumulate mlocked memory until process exit and eventually hit the memlock limit. Attach a deleter/finalizer or update the owning call sites to free these native allocations explicitly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 780f786. NativePinnedMemory.pin() now attaches a weakref.finalize to the returned tensor that frees the page-locked allocation (free_cpu_locked_tensor) when the tensor is garbage-collected — matching torch.pin_memory()'s lifetime, so callers that never call unpin_memory() no longer accumulate mlocked host memory. An explicit unpin() cancels the finalizer (.detach()) to avoid a redundant free. Added test_pin_freed_on_gc (drops the tensor, forces GC, asserts the range is released); the full pin-memory suite is 22/22 on a GPU node.
| # None so the accelerator uses its (torch) pinning path. When native is | ||
| # selected but the async-io ops cannot be built, constructing the manager | ||
| # raises rather than silently falling back. | ||
| if os.environ.get("DS_PIN_MEMORY_BACKEND", "torch").strip().lower() != "native": |
There was a problem hiding this comment.
This introduces the user-visible DS_PIN_MEMORY_BACKEND=native mode, but the commit does not add any documentation for enabling it, the AIO build requirement, or the new explicit unpin behavior. The project rules require documentation updates for new features, so users otherwise have to infer these operational requirements from code and tests.
AGENTS.md reference: AGENTS.md:L26-L26
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Documented under RTD Memory Usage as Host Memory Pinning (docs/code-docs/source/memory.rst):
DS_PIN_MEMORY_BACKEND=torch|nativeselection and defaults- async-io build / memlock requirements for
native - accelerator
pin_memory/is_pinned/unpin_memoryAPIs - native lifetime (
weakreffinalizer + explicit unpin in ZeRO/destroy())
Also refreshed the older Pinned Memory accounting notes to the current offload_*.pin_memory config knobs and cross-linked them to this section.
|
Addressed Codex review (r3713658448) by landing #8212 first, then rebasing this PR onto updated With the process-wide shared |
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 <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
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 <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
|
Addressed the native pin lifetime concern with two complementary paths:
Coverage:
|
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 <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Make Host Memory Pinning a peer of Memory Requirements, with Backend Selection, Requirements for native, and Lifetime and unpinning nested under it. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
|
@delock , @stas00, @tohtana see this page for docs preview of this feature: |
Colocate NativePinnedMemory unit tests with the destroy-path pin_memory suite under the v1 tree. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
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 DeepNVMe async-io (mlock) allocator instead oftorch.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.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 the async-io op 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.compilepaths route throughget_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 / ZenFlowdestroy()explicitly unpins optimizer-owned CPU-offload buffers under the native backend.docs/code-docs/source/memory.rst); also pushed tortd-staging.tests/unit/v1.Test plan
pre-commit(yapf, flake8, check-license, check-torchdist, codespell) passes on all changed files.masterafter 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.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.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).tests/unit/v1/suite on a GPU node — 1134 passed, 63 skipped, exit 0 (~2.5h).DS_PIN_MEMORY_BACKEND=native, a buffer fromget_accelerator().pin_memory(and a narrow of it) isis_pinnedon a different AIO handle than the allocator.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