Split native host pinning into standalone pin_memory op by sfc-gh-truwase · Pull Request #8236 · deepspeedai/DeepSpeed · GitHub

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

tests/unit/v1/pin_memory/test_pin_memory_op.py (pin without async_io)
tests/unit/v1/nvme/test_pinned_manager.py cross-op recognition (pin_handle ↔ aio_handle)
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)
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

Native page-locked allocations (posix_memalign/mlock) no longer require libaio or AIO worker threads. Compile the pin manager only in pin_memory and share it with async_io/gds so bounce-buffer recognition stays process-wide. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 38e2ea6bcf

ℹ️ 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".

Precompile pin_memory with async_io

When a wheel is built with DS_BUILD_AIO=1 (or DS_BUILD_GDS=1) while DS_BUILD_PIN_MEMORY remains at the default false, setup.py still records async_io as installed but pin_memory as not installed. This unconditional dependency load then routes through OpBuilder.load for pin_memory and JIT-compiles it at runtime before loading the already-precompiled AIO op, so deployments without ninja/a compiler fail even though they precompiled AIO; either force pin_memory into the installed ops whenever AIO/GDS is precompiled or avoid requiring a runtime JIT here.

Useful? React with 👍 / 👎.

Addressed in 0dfc557: setup.py now precompiles pin_memory whenever async_io/gds are precompiled (even if DS_BUILD_PIN_MEMORY is left at the default), so AIO/GDS wheels do not require a runtime JIT of the shared pin manager.

Guard the POSIX pin_memory op on Windows

When building wheels on Windows, setup.py defaults to prebuilding compatible ops, but this new op inherits the default compatibility and its source uses POSIX-only headers/APIs such as sys/mman.h, unistd.h, posix_memalign, sysconf, and mlock. A Windows build will now try to compile this source and fail during installation, so add a Windows implementation or make PinMemoryBuilder.is_compatible() return false on Windows.

Useful? React with 👍 / 👎.

Addressed in 0dfc557: PinMemoryBuilder.is_compatible() returns False on win32 because the allocator uses POSIX mlock/posix_memalign.

Document the new pin_memory API

This commit exposes a new public deepspeed.ops.pin_memory / PinMemoryBuilder / pin_handle path and changes how DeepNVMe pinning is shared, but the changed files include no documentation or tutorial updates explaining how users should call it or how it differs from the existing PyTorch/accelerator pinning paths. The repository instructions require documentation updates for new features.

AGENTS.md reference: AGENTS.md:L26-L26

Useful? React with 👍 / 👎.

Addressed in 0dfc557 (DeepNVMe tutorial) and follow-ups in docs/code-docs/source/memory.rst (Host Memory Pinning): documents PinMemoryBuilder / pin_handle, shared process-wide manager with AIO/GDS, and how it differs from torch/accelerator pinning.

…d docs - PinMemoryBuilder.is_compatible() returns False on Windows (POSIX-only mlock/posix_memalign). - setup.py precompiles pin_memory whenever async_io/gds are precompiled so wheels without a runtime compiler don't hit a JIT build for the shared manager. - Document the standalone pin_memory op / pin_handle API in the DeepNVMe tutorial. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Addressed the Codex review feedback in 0dfc557:

P1 (runtime JIT of pin_memory): setup.py now precompiles pin_memory whenever async_io/gds are precompiled, so wheels built with DS_BUILD_AIO=1/DS_BUILD_GDS=1 (and default DS_BUILD_PIN_MEMORY) don't require a compiler at runtime for the shared pin manager.
P2 (Windows): PinMemoryBuilder.is_compatible() returns False on win32 (allocator uses POSIX mlock/posix_memalign).
P2 (docs): Documented the standalone pin_memory op / pin_handle API in the DeepNVMe tutorial.

Add a Host Memory Pinning section to memory.rst covering torch vs the standalone pin_memory op / pin_handle API, shared manager with AIO/GDS, and build/memlock requirements. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Describe the pin_memory op without framing it against libaio or async_io. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Align wording with rtd-staging; this PR only documents torch/accelerator pinning and the standalone pin_memory op. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Microsoft header should be removed

Removed in a6f6ab6 (SPDX + DeepSpeed Team only).

Same, should remove Microsoft header

Removed in a6f6ab6 (and across the other new pin_memory sources).

Why in CPU accelerator cannot just use PinMemoryBuilder defined in op_builder/pin_memory.py

Agreed — host pinning is accelerator-agnostic. op_builder/cpu/pin_memory.py (and the xpu/npu/supa stubs) now just re-export op_builder.pin_memory.PinMemoryBuilder so there is a single implementation. Kept the thin module so the CPU accelerator whitelist / from .pin_memory import PinMemoryBuilder path still resolves.

…der on CPU New pin_memory sources use SPDX + DeepSpeed Team only. CPU (and other accelerator) builders re-export op_builder.pin_memory.PinMemoryBuilder instead of duplicating the class. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Codex + test plan follow-up

Codex (0dfc557, already on branch):

P1: setup.py precompiles pin_memory whenever async_io/gds are precompiled
P2: PinMemoryBuilder.is_compatible() is False on Windows
P2: Docs in DeepNVMe tutorial + docs/code-docs/source/memory.rst

Test plan (HEAD a6f6ab6e):

AIO/GDS/pin_memory unit tests: 131 passed (autorun job 20260809T123310Z)
ds_report: pin_memory ............. [NO] ....... [OKAY]
Without libaio: pin_memory stays compatible when io_submit is mocked missing; async_io does not (20260809T124859Z)

NativePinnedMemory now uses PinMemoryBuilder/pin_handle instead of AsyncIOBuilder, matching the #8236 split so native mode no longer requires libaio. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

…eepspeedai#8211) ## Summary Adds a native host-memory pinning backend, selectable via the `DS_PIN_MEMORY_BACKEND` environment variable (defaults to `torch`). When set to `native`, CPU memory is page-locked through the standalone DeepSpeed `pin_memory` op (`PinMemoryBuilder` / `pin_handle`, `posix_memalign` + `mlock`) instead of `torch.pin_memory()`. Stacked on deepspeedai#8236 (standalone `pin_memory` op, now on `master`). Native allocations go through `pin_handle`, so DeepNVMe I/O handles recognize them via the process-wide manager and skip bounce buffers — without requiring libaio / AIO worker threads. - **New `deepspeed/utils/pin_memory.py`**: a process-wide shared `NativePinnedMemory` manager 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. It fails early with a clear error if the `pin_memory` op cannot be built (no silent torch fallback). Native pins also use a `weakref` finalizer so GC releases mlocked pages when tensors are dropped without an explicit unpin. - **Accelerator owns dispatch**: `pin_memory` drops `align_bytes` and gains `make_copy`/`match_shape`; `is_pinned` is FakeTensor/meta-tensor safe; new `unpin_memory` (native frees, torch no-op). Subclasses retain only the device-specific `_torch_pin_memory`/`_torch_is_pinned` primitives. Preserves master's `track_pinned_memory` accounting (CPU torch no-op still bypasses it). - **Consolidation**: XPU's bespoke `align_bytes=0` path is folded into the shared native backend. - **Callers**: `compile` paths route through `get_accelerator()`. Swap-tensor buffers continue to allocate via I/O handles; with the shared manager they interoperate with native-pinned tensors. ZeRO / ZenFlow `destroy()` explicitly unpins optimizer-owned CPU-offload buffers under the native backend. - **Docs**: Host Memory Pinning section under RTD Memory Usage (`docs/code-docs/source/memory.rst`). - **Tests**: unit tests for the native manager, accelerator pinning APIs, destroy-path unpin, and cross-op recognition with AIO. ## Test plan - [x] Rebased onto `master` after deepspeedai#8236 merge; retargeted `NativePinnedMemory` from `AsyncIOBuilder` → `PinMemoryBuilder`. - [x] `pre-commit` on changed files. - [x] Focused UTs on GPU (`tunji-h200-n1g2-ds2-0`, job `20260809T183311Z`): `tests/unit/v1/pin_memory/` + `tests/unit/v1/accelerator/test_accelerator.py` + `tests/unit/v1/nvme/test_pinned_manager.py` — **30 passed**. - [x] Bounce-buffer / cross-op smoke: under `DS_PIN_MEMORY_BACKEND=native`, a `pin_handle` buffer is `is_pinned` on a separate AIO handle. Made with [Cursor](https://cursor.com) --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

sfc-gh-truwase requested review from loadams, tjruwase and tohtana as code owners August 8, 2026 23:08

sfc-gh-truwase mentioned this pull request Aug 8, 2026

Add native (DeepNVMe) host-memory pinning backend for accelerators #8211

Merged

4 tasks

chatgpt-codex-connector Bot reviewed Aug 8, 2026

View reviewed changes

sfc-gh-truwase and others added 3 commits August 8, 2026 23:48

Docs: drop libaio/async_io from pin_memory op RTD

e433994

Describe the pin_memory op without framing it against libaio or async_io. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Docs: clarify Host Memory Pinning cross-ref (no native backend)

a043294

Align wording with rtd-staging; this PR only documents torch/accelerator pinning and the standalone pin_memory op. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

sfc-gh-truwase requested review from delock and removed request for loadams August 9, 2026 03:33

delock reviewed Aug 9, 2026

View reviewed changes

delock approved these changes Aug 9, 2026

View reviewed changes

sfc-gh-truwase added this pull request to the merge queue Aug 9, 2026

Merged via the queue into master with commit 5ad9a97 Aug 9, 2026
14 checks passed

sfc-gh-truwase deleted the tjruwase/pin-memory-op branch August 9, 2026 13:50

+                      # Pin manager is compiled only into pin_memory; load it first so aio can
+                      # resolve the shared manager (and related symbols) across .so boundaries.
+                      from .pin_memory import PinMemoryBuilder
+                      PinMemoryBuilder().load(verbose=verbose)


		# DeepSpeed Team

		from ..op_builder.pin_memory import PinMemoryBuilder

		from deepspeed.ops.op_builder.pin_memory_load import load_pin_memory_module


		class PinMemoryBuilder(CPUOpBuilder):

		#include <sys/mman.h>
		#include <unistd.h>

		@@ -0,0 +1,43 @@
		// Copyright (c) Microsoft Corporation.

		@@ -0,0 +1,35 @@
		// Copyright (c) Microsoft Corporation.

Split native host pinning into standalone pin_memory op by sfc-gh-truwase · Pull Request #8236 · deepspeedai/DeepSpeed · GitHub

Uh oh!

Split native host pinning into standalone pin_memory op - #8236

Split native host pinning into standalone pin_memory op#8236sfc-gh-truwase merged 6 commits intomasterdeepspeedai/DeepSpeed:masterfrom tjruwase/pin-memory-opdeepspeedai/DeepSpeed:tjruwase/pin-memory-opCopy head branch name to clipboard

Conversation

sfc-gh-truwase commented Aug 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Follow-up

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Aug 8, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Aug 8, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Aug 8, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-truwase commented Aug 8, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sfc-gh-truwase commented Aug 9, 2026

Codex + test plan follow-up

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Split native host pinning into standalone pin_memory op#8236
sfc-gh-truwase merged 6 commits into
masterfrom
tjruwase/pin-memory-op

sfc-gh-truwase commented Aug 8, 2026 •
edited
Loading