Share DeepNVMe pinned-tensor manager and route swap buffers through I/O handles - #8212
Conversation
…/O handles DeepNVMe skips the bounce-buffer copy only when a buffer is torch-pinned or managed by the handle's pinned-tensor manager. That manager was per-handle and recognized only exact base pointers, so buffers allocated by one handle (or narrows/views of a shared pool submitted through a different read/write handle) were not recognized and always bounced. - C++: make deepspeed_pin_tensor_t a process-wide shared() singleton guarded by a mutex, and make is_managed range-based so slices/views of a locked buffer are recognized. Switch deepspeed_io_handle_t and cpu_op_desc_t to hold the manager via shared_ptr, and add handle.is_pinned() with aio/gds bindings. - Python: the swap subsystem now obtains pinned host memory via new_cpu_locked_tensor and queries pinned status via handle.is_pinned instead of the accelerator. SwapBufferManager/SwapBufferPool and the optimizer swappers take an aio_handle; optimizer-swapper subclasses create the handle before super().__init__ so it can be threaded through. - Add a focused nvme test for narrow recognition and cross-handle sharing. 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: dd2d494965
ℹ️ 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".
- handle.is_pinned() now returns is_pinned() || is_managed(), mirroring cpu_op_desc_t's direct-I/O eligibility check. Reporting only is_managed() misclassified torch-pinned destinations (e.g. ZeRO-3 fp16 flat CPU memory) as unpinned, needlessly forcing the swap-buffer path and risking pool exhaustion. - SwapBufferManager now releases its page-locked buffers via free_cpu_locked_tensor() on destruction. With the process-wide pinned-tensor manager, dropping a manager no longer reclaims its buffers, so repeatedly building NVMe swappers would otherwise accumulate locked memory. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
|
After this PR, no caller passes align_bytes=0 in |
is_managed() only checked that a buffer's start pointer fell inside a locked region, so a buffer beginning inside a region but ending past it (an unpinned tail) was misreported as managed. Also require the buffer's end (ptr + nbytes) to fall within the same region. Legitimate swap narrows/views are unaffected since they lie fully within their allocation. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
@delock yes please help with the cleanup. FYI, I am working to promote DeepSpeed pin memory to a first class utility in here. I think this is useful because of the tangible differences from torch version. Additionally, we shall make two changes to DeepSpeed code base
I will like to discuss more on these changes. Thanks! |
2 participants
Summary
DeepNVMe skips the bounce-buffer copy only when a buffer is torch-pinned or managed by the handle's pinned-tensor manager. That manager was per-handle and recognized only exact base pointers, so buffers allocated by one handle — or narrows/views of a shared pool submitted through a different read/write handle — were not recognized and always bounced.
This PR makes the pinned-tensor manager a process-wide shared instance with range-based recognition, and switches the swap subsystem to obtain pinned memory and query pinned status through its I/O handles.
Changes
deepspeed_pin_tensor_tis now a process-wideshared()singleton guarded by astd::mutex;is_managedis range-based so slices/views of a locked buffer are recognized.deepspeed_io_handle_tandcpu_op_desc_thold the manager viastd::shared_ptr.handle.is_pinned(buffer)with bindings inpy_ds_aio.cppandpy_ds_gds.cpp(GDS inherits the base).SwapBufferManager/SwapBufferPooland the optimizer swappers take anaio_handle; buffers are allocated vianew_cpu_locked_tensorand pinned status is queried viahandle.is_pinned.super().__init__so it can be threaded through.tests/unit/v1/nvme/test_pinned_manager.pycovers narrow/view recognition and cross-handle sharing.Test plan
pre-commit(yapf/flake8/clang-format/check-license) on all touched files.tests/unit/v1/nvme/test_pinned_manager.py— 3/3 pass (narrow recognition, cross-handle sharing, unmanaged buffer).tests/unit/v1/nvme/+tests/unit/utils/test_pin_memory.py+tests/unit/v1/accelerator/test_accelerator.py— 146 pass.tests/unit/runtime/zero/test_nvme_checkpointing.py): reproduces the pre-existing baseline exactly (no regression; the failing optimizer-on-NVMe configs fail identically onmaster).Made with Cursor