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 1 commit into
Conversation
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>
![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 👍 / 👎.
1 participant
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().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).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.swap_tensorandcompilepaths route throughget_accelerator().tests/unit/v1.Test plan
pre-commit(yapf, flake8, check-license, check-torchdist, codespell) passes on all changed files.tests/unit/utils/test_pin_memory.pyandtests/unit/v1/accelerator/test_accelerator.pypass on a GPU node (18 passed), exercising both the torch and native paths plus the fail-early behavior when async-io is unavailable._torch_*primitives.Made with Cursor