Route ZeRO/SuperOffload pin sites through accelerator pin_memory

[sfc-gh-truwase]

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.

[chatgpt-codex-connector]

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

[delock]

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]

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

[xylian86]

I have one concern on lifetime of these native-pinned buffers. During reload_states(non_blocking=True), the CPU-to-GPU copy may still be running when the HP/LP owner lists are deleted.

[sfc-gh-truwase]

@xylian86, can you explain a bit more? Is your concern specific to this PR or more generally about the lifetime of native pinning?

[xylian86]

Yes, for sure. The flow involves three steps:

Step 1. start copy

DeepSpeed/deepspeed/runtime/zero/stage_1_and_2.py

Line 3019 in aa3914d

buf.data = buf.data.to(device, non_blocking=non_blocking)

Step 2. delete source memory

DeepSpeed/deepspeed/runtime/zero/stage_1_and_2.py

Line 3021 in aa3914d

del self.hp_params_pin_buffers

Step 3. wait for copy to finish

DeepSpeed/deepspeed/runtime/zero/stage_1_and_2.py

Lines 3062 to 3063 in aa3914d

	if non_blocking:
	get_accelerator().synchronize()

With PyTorch’s pinned-memory allocator, dropping the tensor does not immediately release its underlying memory. The allocator tracks the CUDA stream using the buffer and retains the allocation until the copy completes.

My understanding is that DeepSpeed’s native allocator tracks Python object ownership but does not track the CUDA stream using the buffer. Therefore, step 2 may free the source allocation before the synchronization in step 3.

Could we retain these source buffers until after synchronization, or add a native nonblocking GPU test to confirm that the current lifetime is safe?

[sfc-gh-truwase]

Good catch.

1. Yes, we can retain the source buffers until after synchronization
2. I will add native nonblocking GPU tests
3. I will investigate whether Register native pinned host memory with CUDA for GPU DMA #8283 enables CUDA stream tracking

[sfc-gh-truwase]

Thanks @xylian86 — you were right, and the fix is pushed in aebb2819.

1. Retain the source buffers until after synchronization. reload_states now holds every host source in a local list and releases it only after the synchronize():

        pending_host_buffers = []
        ...
            if hasattr(self, "hp_params_pin_buffers"):
                pending_host_buffers.append(self.hp_params_pin_buffers)
                del self.hp_params_pin_buffers
        ...
        if non_blocking:
            get_accelerator().synchronize()
        # The copies have completed, so the host sources can be released.
        del pending_host_buffers

reload_optimizer_states had the same hazard — it replaced state[k], dropping the pinned source right away — so it now returns the replaced host tensors and reload_states keeps them in the same list.

ZeRO-3 needs no change: hp_params_pin_buffers, lp_param_contiguous_pin_buffer, lp_grad_partitions_flat_pin_buffers, and the Adam offload buffers all stay reachable as attributes across reload, so nothing is freed mid-copy.

2. Native non-blocking GPU tests. test_reload_states_holds_pin_buffers_until_sync drives the real reload_states(non_blocking=True) with a native-pinned hp buffer and records whether the allocation is still live at the moment synchronize() is called. On 1x H200 it passes with the fix and fails assert [False] == [True] with the two source files reverted, so it genuinely catches the use-after-free. test_reload_optimizer_states_returns_host_sources covers the helper contract.

Full run on H200:

Run	Result
test_offload_route.py (native)	4 passed
same, fix reverted	fails as expected
test_offload_states.py, 2 GPUs, torch backend	128 passed
test_offload_states.py, 2 GPUs, native backend	128 passed

3. Does #8283 give stream tracking? No. cudaHostRegister only makes the pages DMA-capable; deferred, stream-ordered reuse is a property of torch's caching host allocator, which the native (posix_memalign + mlock) allocator does not have. Registration arguably makes the hazard more visible, since the copy becomes a true async DMA rather than a synchronous staged copy. I documented the caveat in memory.rst so future call sites know that non_blocking=True out of a native-pinned buffer requires holding a reference until synchronization.