DeepSpeed
e1d6b4fe - Share DeepNVMe pinned-tensor manager and route swap buffers through I/O handles (#8212)

Commit
14 hours ago
Share DeepNVMe pinned-tensor manager and route swap buffers through I/O handles (#8212) ## 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 - **C++** - `deepspeed_pin_tensor_t` is now a process-wide `shared()` singleton guarded by a `std::mutex`; `is_managed` is **range-based** so slices/views of a locked buffer are recognized. - `deepspeed_io_handle_t` and `cpu_op_desc_t` hold the manager via `std::shared_ptr`. - Added `handle.is_pinned(buffer)` with bindings in `py_ds_aio.cpp` and `py_ds_gds.cpp` (GDS inherits the base). - **Python (swap subsystem only)** - `SwapBufferManager`/`SwapBufferPool` and the optimizer swappers take an `aio_handle`; buffers are allocated via `new_cpu_locked_tensor` and pinned status is queried via `handle.is_pinned`. - Optimizer-swapper subclasses create their handle before `super().__init__` so it can be threaded through. - **Test**: `tests/unit/v1/nvme/test_pinned_manager.py` covers narrow/view recognition and cross-handle sharing. ## Test plan - [x] `pre-commit` (yapf/flake8/clang-format/check-license) on all touched files. - [x] `tests/unit/v1/nvme/test_pinned_manager.py` — 3/3 pass (narrow recognition, cross-handle sharing, unmanaged buffer). - [x] `tests/unit/v1/nvme/` + `tests/unit/utils/test_pin_memory.py` + `tests/unit/v1/accelerator/test_accelerator.py` — 146 pass. - [x] Swap smoke test (`tests/unit/runtime/zero/test_nvme_checkpointing.py`): reproduces the pre-existing baseline exactly (no regression; the failing optimizer-on-NVMe configs fail identically on `master`). Made with [Cursor](https://cursor.com) --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Parents
Loading