[pengdurice]

ZeRO-3: copy offloaded master weights back to the device through pinned buffers

Part of: #8672 (Q4 2026 roadmap), "Enable training of large scale models".

Summary

With offload_optimizer.device: cpu, ZeRO-3 copies each sub-group's updated float32 master partition back to the
device from host memory that is never pinned, even with pin_memory: true (the default since #8207). This PR copies
it through two fixed-size pinned buffers instead, with the same values.

• Qwen2.5-32B, 8 H200, whole optimizer on the CPU: 11.96 s to 10.35 s per step (-13.5%), same peak memory.
• MiniMax-M3 (all 60 layers, 4 nodes x 8 H200, 65% of the optimizer on the CPU): 126.9 s to 123.9 s per step.
• Extra host memory: a fixed 256 MiB of pinned memory per rank (bf16).

Duplicate check

Searched issues and PRs on 2026-10-13 ("offload pinned fp32 partition", "zero3 offload pin_memory copy",
"fp32_partitioned_groups_flat", "offload ratio pinned", "pageable offload"). Related, not duplicates: #8207 (merged:
pinned the ZeRO-3 gradient offload buffers, default pin_memory: true; master partitions unchanged), #8256 (merged:
routes pin sites through get_accelerator().pin_memory(); this PR uses the same accelerator calls), #8618 (closed:
a CPU-offload benchmark).

Sibling path: ZeRO-1/2 does not have the defect. It pins the float32 partition and casts into a pinned bf16 buffer
of the partition's size before a non_blocking copy (stage_1_and_2.py lines 546-558 and 2841-2844).

Root cause

_create_fp32_partitions (stage3.py) builds each CPU master partition with .to('cpu', dtype=...) and
.clone().detach(), which is pageable memory whatever pin_memory says. After the CPU optimizer step,
_reassign_or_swap_out_partitioned_parameters runs fp16_flat.data.copy_(fp32_flat.data): a blocking copy from CPU
float32 to device bf16. For a blocking copy with a dtype change, PyTorch converts on the CPU into a new pageable bf16
tensor of the whole partition, copies it from pageable memory and synchronizes.

Found in a profile (MiniMax-M3, first 8 layers, 8 H200, 75% of the optimizer on the CPU, rank 0): 13 bare
aten::copy_ host-to-device copies of 576 MiB per step from pageable memory, at 3-7 GB/s, 1,648 ms per step.

The change (deepspeed/runtime/zero/stage3.py, 49 added, 2 removed lines)

Where	What
_reassign_or_swap_out_partitioned_parameters	uses the new copy when offload_optimizer.pin_memory is true, the source is on the CPU and not pinned (get_accelerator().is_pinned), and the destination is on the device; every other case keeps dst.copy_(src)
new _copy_through_pinned_staging(dst, src)	two pinned buffers of 2**26 elements in dst's dtype (get_accelerator().pin_empty, created on first use). Per piece: wait for the previous copy out of this buffer (an event), cast the piece into the buffer on the CPU, copy_(..., non_blocking=True) to the device, record an event. Before returning, wait for the last piece
__init__, _unpin_offload_buffers	the buffer size and state; release of the buffers (matters under DS_PIN_MEMORY_BACKEND=native)

Design notes:

• Why not pin the master partitions, as ZeRO-1/2 does: a non_blocking copy from pinned float32 into a bf16
  tensor converts on the device and needs a float32 device temporary of the sub-group's size. Avoiding that needs a
  pinned bf16 buffer of each partition's size, as ZeRO-1/2 keeps. Either way pinned host memory would grow with the
  model, and PyTorch's pinned allocator rounds each allocation up to a power of two. The two buffers here have a fixed
  size; 2**26 is a power of two, so nothing is rounded up.
• Same values: the CPU cast is the one the blocking copy did before (round to nearest even), so the device
  partition is bitwise the same.
• Stream order: the pieces are copied on the current stream, and an event guards each buffer's reuse. The final
  wait keeps what callers had before: the partition is on the device when the function returns. The next forward's
  parameter all-gathers run on their own stream, and it does not wait for the main stream.
• Host synchronization: at most one piece-sized wait per buffer reuse and one at the end. The old path blocked
  for the whole copy.

Measured

Qwen2.5-32B (all 64 layers, real weights), 1 node of 8 H200 (143,771 MiB, driver 580.126.09), torch 2.8.0+cu126.
Settings:

• ZeRO-3 with offload_optimizer: {device: cpu, pin_memory: true};
• float32 master weights and Adam states, DeepSpeedCPUAdam with 12 threads per rank;
• bf16, one 8,192-token PG-19 sequence per GPU, activation checkpointing, lr 1e-6;
• mean of steps 3-8; one run per row.

	master dcd8595c2	this PR
step time	11.96 s	10.35 s
step time with wall_clock_breakdown	11.72 s	10.01 s
rank 0: forward / backward	1.21 / 3.73 s	1.21 / 3.72 s
rank 0: step / of which the sub-group loop (optimizer_step)	5.98 / 5.73 s	5.03 / 3.64 s
peak device memory (rank 0)	34.66 GB	34.66 GB
loss at steps 1 / 2 / 8	0.022649 / 0.976567 / 1.746324	0.022649 / 0.976611 / 1.746344

The losses differ from step 2 by run-to-run variation. Two master runs differ by as much: step 2 is 0.976376 with
wall_clock_breakdown and 0.976567 without it. The flash-attention backward adds with atomics. The bitwise check is
the unit test below.

MiniMax-M3 runs (DeepSpeed master + unmerged MiniMax-M3 model code: #8658, #8652)

Settings for both: ZeRO-3 + AutoEP (expert parallel 8), real weights, each 65,536-token sequence split over 4 GPUs,
65% of the optimizer on the CPU (offload_optimizer.ratio 0.65), float32 master + Adam.

• 4 nodes x 8 H200, all 60 layers, same nodes, back to back: 126.9 s without, 123.9 s with
  (mean of steps 3-10); peak device memory 132.4 GB in both.
• 1 node, first 8 layers: 14.09 s without, 12.46 s with; losses equal at every printed step.

This change, and the fused CPU Adam change stacked on it, were in the DeepSpeed tree of our MiniMax-M3 and GLM-5.2
runs against Megatron-LM (real weights, 65,536-token sequences, 80-95% of the optimizer on the CPU):

Model	Nodes	DeepSpeed, s/step	Megatron-LM, s/step	How measured
MiniMax-M3, all 60 layers	4 x 8 H200	85.24	119.35	one job, same nodes, mean of steps 4-23
GLM-5.2, first 38 layers	4 x 8 H200	62.03	63.73	one job, same nodes, mean of steps 4-23
GLM-5.2, all 78 layers	8 x 8 H200	130.58	135.56	separate jobs on 7 of the same 8 nodes, mean of steps 6-10

On one node (first 4 layers of each model), master plus this PR, the fused CPU Adam PR and the two model
PRs (#8658, #8659) gives the same loss at every step as that tree.

Tests

New: TestZero3OffloadPinnedCopyBack in tests/unit/v1/zero/test_zero.py (2 ranks, ratio 1.0 and 0.5,
sub_group_size 20,000 so partial offload has sub-groups on both sides).

• It trains 3 steps with pin_memory: false (direct copy) and with pin_memory: true, and requires the gathered
  parameters to be bitwise equal.
• It also requires that only the pinned run used the staging buffers. Otherwise the test would compare the direct
  copy with itself.
• The piece size is set to 1,000 elements, so each step refills both buffers many times.

The test is in tests/unit/v1 so that PR CI selects it. With this diff, the modal-torch-latest selector
(ci/tests_fetcher.py) picks 85 test files, including tests/unit/v1/zero/test_zero.py. I simulated the selector's
diff with the working-tree files, since nothing was committed.

Commands, from tests/, on 1 node of 8 H200 (143,771 MiB, driver 580.126.09), torch 2.8.0+cu126, pytest 9.1.1:

# master + only the new test: 2 failed, AttributeError: no attribute '_staging_buffers'
# master + this PR:
pytest unit/v1/zero/test_zero.py -k TestZero3OffloadPinnedCopyBack        # 2 passed
pytest unit/v1/zero/test_zero.py unit/v1/zero/test_zero_user_backward.py unit/v1/pin_memory/ \
    -k "offload or Offload or pin or Pin"                                  # 61 passed, 1 skipped (torch_npu)
pytest unit/v1/zero/test_offload_states.py                                 # 128 passed
pytest unit/runtime/zero/test_zero_offloadpp.py unit/v1/pin_memory/        # 50 passed
# on master 5f266e26d (2026-10-15, this change plus the fused CPU Adam PR stacked on it):
pytest unit/v1/zero/test_zero.py::TestZero3OffloadPinnedCopyBack            # 2 passed
pytest unit/runtime/zero/test_zero_offloadpp.py                             # 3 passed (master: 3 passed)

The last two ran on the first version of this patch, whose stage3.py differs from this one only in two comments;
the timings above also used the first version.

Compatibility

Cell	Validated
ZeRO-3, CPU optimizer offload, pin_memory true, whole optimizer on the CPU	yes: new test (ratio 1.0); Qwen2.5-32B
ZeRO-3, partial offload (ratio < 1)	yes: new test (ratio 0.5); MiniMax-M3 on 1 and 4 nodes
ZeRO-3, pin_memory false	unchanged path; the new test's reference
bf16	yes. fp16: not run (same path; the cast is the same CPU cast)
ZeRO-3, NVMe optimizer offload	not run. Expected unchanged: the swapped-in buffers are pinned, so the direct copy is kept
ZeRO-3 + ZenFlow, DS_PIN_MEMORY_BACKEND=native	not run
ZeRO-1/2	not affected
Accelerators other than CUDA	not run; only get_accelerator() calls are used (pin_empty, is_pinned, Event, unpin_memory)

No config, checkpoint or public API change.

AI assistance

This change, its tests and this description were prepared with AI assistance (Claude).

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-20T21:11:59.347603Z	f8d5502	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[delock]

suppose we run v1 tests on CPU, then these tests should be skipped on CPU because offload is not a CPU scenario.

[pengdurice]

great catch, added a line to skip cpu below (similar to what TestZero3DictFwd does it);-)

[delock]

Hi @pengdurice , thanks for your PR. I have left my comments.

[sfc-gh-truwase]

@pengdurice can you confirm why self.fp16_partitioned_groups_flat and self.fp32_partitioned_groups_flat are not pinned at construction if pin_memory option is selected

[pengdurice]

Hi @sfc-gh-truwase , thank you for the review. For self.fp16_partitioned_groups_flat, it is on the GPU unless offloading parameters are on, so pinning doesn't apply here. For self.fp32_partitioned_groups_flat, it is indeed pinned, but during copying back to the GPU, it uses an pageable memory by PyTorch which is the inefficient part. So, this PR creates two fixed sized pinned buffer. LMK if that's clear. Thank you!