Route ZeRO/SuperOffload pin sites through accelerator pin_memory by sfc-gh-truwase · Pull Request #8256 · deepspeedai/DeepSpeed

sfc-gh-truwase · 2026-08-15T23:01:29Z

Summary

Route ZeRO-1/2 offload_states hp/lp pins, offload_optimizer_states, ZeRO++ secondary shards, and SuperOffload grad buffers through get_accelerator().pin_memory() so DS_PIN_MEMORY_BACKEND=native applies.
Use make_copy=False for empty scratch destinations to avoid an extra host alloc/copy under the native backend.
Clarify in memory.rst that ZeRO offload_*.pin_memory (whether) is orthogonal to DS_PIN_MEMORY_BACKEND (how); env-only, no ds_config backend field.
Add tests/unit/v1/pin_memory/test_offload_route.py covering offload_optimizer_states under native.

Test plan

pre-commit run --files on touched paths (already run locally)
DS_PIN_MEMORY_BACKEND=native + pytest tests/unit/v1/pin_memory/test_offload_route.py (+ pin_memory/accelerator UTs) on H200 host with pin_memory op — 23 passed (incl. both offload_route tests)
Smoke ZeRO-1/2 CPU offload with default torch backend — test_offload_states.py + test_destroy_unpin.py — 133 passed
SuperOffload worker pin site: source uses get_accelerator().pin_memory(..., make_copy=False); native empty-buffer pin/unpin smoke OK

Evidence: autorun job-20260818T143819Z on e6dc5f9d / tjruwase/pin-memory-route-zero (EXIT 0). GitHub CI also green on the tip commit.

Honor DS_PIN_MEMORY_BACKEND for core CPU offload buffers and avoid an extra host copy on empty scratch pins via make_copy=False. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

chatgpt-codex-connector · 2026-08-15T23:01:35Z

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

Skip GPU->pinned offload helper coverage on CPUAccelerator and add a native pin-pattern smoke that cpu-torch-latest can run. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

delock · 2026-08-20T15:41:46Z

Hi @sfc-gh-truwase , this PR looks good to me. There is only one thoughts, the phrase " get_accelerator().pin_memory(torch.empty(...), make_copy=False)" keeps repeating and it looks like an idiom. Maybe it worth a place in utility.

sfc-gh-truwase · 2026-08-20T15:58:28Z

@delock great point. I will address in a separate PR.

The native pin backend has no CUDA stream tracking, so dropping a host buffer while its non-blocking copy is in flight is a use-after-free. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

The torch path pinned a copy even for empty scratch destinations, so every make_copy=False site briefly needed two full-size host buffers plus a copy. Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

tohtana

Hi @sfc-gh-truwase, I left a few comments.
These issues are not very critical, but it would be great if we could fix the memory leak one in this PR (though it is a preexisting issue).

tohtana · 2026-08-24T18:13:40Z


    # Pre-allocate reusable pinned memory buffer for gradients
-    pinned_grad_buffer = torch.empty(max_grad_numel, dtype=torch.float32, device='cpu', pin_memory=True)
+    pinned_grad_buffer = get_accelerator().pin_memory(torch.empty(max_grad_numel, dtype=torch.float32, device='cpu'),


Do we want to disable pinning based on offload_optimizer.pin_memory? We might not, but it would be good to raise an error if pin_memory is False.

Good question.

@xylian86 what do you think? Should we condition memory pinning on offload_optimizer.pin_memory or not?

@sfc-gh-truwase Since pinning appears to be a performance optimization rather than a correctness requirement. My preference is to honor offload_optimizer.pin_memory. When it is true, we should keep the current pinned allocation; when it is false, the worker should allocate a regular CPU buffer.

Done in 658dc247. SuperOffload now honors offload_optimizer.pin_memory:

true (default): same pinned scratch buffer as before, via get_accelerator().pin_memory(..., make_copy=False)

false: regular torch.empty(..., device="cpu")

The flag is passed from self.offload_optimizer_pin_memory into SuperOffloadCPUOptimizer and the worker. On worker shutdown we unpin_memory only when we actually pinned.

Tests: test_superoffload_grad_buffer_unpinned_when_disabled and test_superoffload_grad_buffer_pinned_when_enabled (native) both passed locally.

tohtana · 2026-08-24T18:43:16Z

                if not hasattr(self, "hp_params_pin_buffers"):
                    self.hp_params_pin_buffers = [
-                        torch.empty_like(t, device=device).pin_memory() for t in self.single_partition_of_fp32_groups
+                        get_accelerator().pin_memory(torch.empty_like(t, device=device), make_copy=False)


When we use engine.offload_states(pin_memory=True) and reach engine.destroy() without calling reload_states(), engine.destroy() doesn't release the pinned buffers (hp_params_pin_buffers, lp_params_pin_buffers, and self.optimizer.state[*][key]).
This is because we have the condition in destroy:

if not (self.cpu_offload and self.cpu_offload_pin_memory): return

I think this is a preexisting issue but it is now related to NativePinnedMemory manager as well. Can we resolve this issue in this PR?

Fixed in 05ddc1a41. You were right: _unpin_offload_buffers returned immediately unless cpu_offload and cpu_offload_pin_memory, so offload_states(pin_memory=True) then destroy() without reload_states() leaked the host buffers.

destroy() now unpins hp_params_pin_buffers, lp_params_pin_buffers, and any CPU tensors in optimizer.state that is_pinned before that config check. ZeRO-3 already released its contiguous pin buffers unconditionally; it now uses the same helper for the leftover offload_adam_states aliases.

test_native_destroy_frees_offload_states_pins covers stages 2 and 3 with no ZeRO CPU offload configured. H200 job-20260824T224008Z:

Run Result

test_destroy_unpin.py with the fix 7 passed (both new stage-2 and stage-3 cases)

same test, sources reverted to 598ad17a fails offloaded=4 after=4 on both stages

accelerator + pin_memory + offload-route UTs 36 passed

offload_states(pin_memory=True) page-locks host buffers with no regard for the ZeRO CPU-offload config, so destroy() leaked them whenever reload_states() never ran. Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Pinning is a DMA optimization, not a correctness requirement, so a False config should allocate a regular CPU scratch buffer instead of always page-locking. Signed-off-by: tunji-ruwase_snow <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 15, 2026 23:01

Fix CPU CI skip for offload pin-route unit test

e6dc5f9

Skip GPU->pinned offload helper coverage on CPUAccelerator and add a native pin-pattern smoke that cpu-torch-latest can run. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>

sfc-gh-truwase requested review from delock and xylian86 August 19, 2026 14:48

delock approved these changes Aug 20, 2026

View reviewed changes

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

xylian86 reviewed Aug 20, 2026

View reviewed changes

Comment thread deepspeed/runtime/zero/stage_1_and_2.py

sfc-gh-truwase removed this pull request from the merge queue due to a manual request Aug 20, 2026

xylian86 reviewed Aug 24, 2026

View reviewed changes

Comment thread deepspeed/runtime/zero/offload_states.py

sfc-gh-truwase and others added 2 commits August 24, 2026 10:15

Merge branch 'master' into tjruwase/pin-memory-route-zero

598ad17

xylian86 approved these changes Aug 24, 2026

View reviewed changes

tohtana requested changes Aug 24, 2026

View reviewed changes

sfc-gh-truwase and others added 2 commits August 24, 2026 22:30

Navigation Menu

Uh oh!

Route ZeRO/SuperOffload pin sites through accelerator pin_memory - #8256

Route ZeRO/SuperOffload pin sites through accelerator pin_memory#8256
sfc-gh-truwase wants to merge 7 commits into
masterfrom
tjruwase/pin-memory-route-zero

sfc-gh-truwase commented Aug 15, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot commented Aug 15, 2026

Uh oh!

delock commented Aug 20, 2026

Uh oh!

sfc-gh-truwase commented Aug 20, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tohtana left a comment

Uh oh!

tohtana Aug 24, 2026

Uh oh!

sfc-gh-truwase Aug 24, 2026

Uh oh!

xylian86 Aug 25, 2026

Uh oh!

sfc-gh-truwase Aug 25, 2026

Uh oh!

xylian86 Aug 25, 2026

Uh oh!

tohtana Aug 24, 2026

Uh oh!

sfc-gh-truwase Aug 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Run	Result
`test_destroy_unpin.py` with the fix	7 passed (both new stage-2 and stage-3 cases)
same test, sources reverted to `598ad17a`	fails `offloaded=4 after=4` on both stages
accelerator + pin_memory + offload-route UTs	36 passed

Uh oh!

Conversation

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

Uh oh!

Summary

Test plan

Uh oh!

chatgpt-codex-connector Bot commented Aug 15, 2026

Uh oh!

delock commented Aug 20, 2026

Uh oh!

sfc-gh-truwase commented Aug 20, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tohtana left a comment

Choose a reason for hiding this comment

Uh oh!

tohtana Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

sfc-gh-truwase Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

xylian86 Aug 25, 2026

Choose a reason for hiding this comment

Uh oh!

sfc-gh-truwase Aug 25, 2026

Choose a reason for hiding this comment

Uh oh!

xylian86 Aug 25, 2026

Choose a reason for hiding this comment

Uh oh!

tohtana Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

sfc-gh-truwase Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

sfc-gh-truwase commented Aug 15, 2026 •

edited

Loading