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

[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).