DeepSpeed
5ad9a978 - Split native host pinning into standalone pin_memory op (#8236)

Commit
89 days ago
Split native host pinning into standalone pin_memory op (#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`, #8211) will be stacked on this PR once it lands. Made with [Cursor](https://cursor.com) --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Parents
Loading