[pengdurice]

ZeRO-3: pin the offloaded master weights and cast them on the device when copying them back

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 as bf16 after every step.

• pin_memory: true (the default since fix(zero3): async grad offload + pinned offload buffers by default #8207) pins the float32 gradient partitions. It does not pin the master
  partitions, so this copy reads pageable host memory.
• The copy is also a cast. A blocking copy_ from float32 into bf16 casts on the CPU into a new pageable tensor and
  copies from that at a few GB/s.

This PR follows ZeRO-1/2 and #5301:

• with pin_memory: true, each CPU master partition is pinned when it is created, as ZeRO-1/2 does;
• the copy-back moves the pinned float32 master to the device and casts it there, in pieces of 2**26 elements, so the
  device temporary is one piece (0.25 GiB) instead of a whole sub-group.

Results:

• Qwen2.5-32B, 8 H200, whole optimizer on the CPU: 12.08 s to 9.42 s per step (-22.0%), same peak device memory,
  same losses (the new unit test compares the parameters bitwise).
• Host memory: the master partitions are already in host memory and become pinned. PyTorch's pinned allocator rounds
  each allocation up to a power of two: a 1e9-element sub-group (3.73 GiB in float32) reserves 4.00 GiB.

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 pins its master partitions
(stage_1_and_2.py#L563-L575).
#5301 added the pinning and copied the pinned master to the device before the cast; ZeRO-1/2 now casts into a pinned
bf16 buffer of each partition's size instead
(#L3093-L3096).
This PR uses #5301's form, which needs no second pinned buffer.

Root cause

What offload_optimizer.pin_memory: true pins in ZeRO-3 on master, checked at runtime on the earlier version of this
PR, which pins the same buffers as master (one GPU, bf16, DeepSpeedCPUAdam, after deepspeed.initialize and after one
step):

Buffer	Where	Pinned
fp16_partitioned_groups_flat (bf16 parameter partitions)	device	(device memory)
fp32_partitioned_groups_flat[i].grad (float32 gradient partitions)	host	yes (stage3.py#L1339-L1340)
fp32_partitioned_groups_flat (float32 master partitions)	host	no
Adam exp_avg, exp_avg_sq	host	no

With offload_param, fp16_partitioned_groups_flat is on the host and already pinned: it is a view of
param_groups_fp16_flat_cpu_memory (stage3.py#L906).

_create_fp32_partitions builds each CPU master partition with .to('cpu', dtype=...) and .clone().detach()
(stage3.py#L1157-L1160),
which gives pageable memory. 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, which PyTorch does by casting
on the CPU into a new pageable bf16 tensor of the whole partition, copying it from pageable memory and synchronizing.

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, 43 added, 4 removed lines)

Where	What
_create_fp32_partitions	with offload_optimizer.pin_memory, pins each CPU master partition after clone().detach() (get_accelerator().pin_memory)
_reassign_or_swap_out_partitioned_parameters	uses the new copy for a master pinned there whose destination is on the device; every other case keeps dst.copy_(src): pin_memory false, sub-groups kept on the device (ratio < 1), NVMe-swapped sub-groups
new _copy_pinned_master_to_device(dst, src)	per piece of 2**26 elements: src[piece].to(device, non_blocking=True), then dst[piece].copy_(...), which casts on the device; at the end, records an event and waits for it
new _master_pinned_at_construction(sub_group_id)	the condition above (pinned, on the CPU, not NVMe-swapped); also used by _unpin_offload_buffers
__init__, _unpin_offload_buffers	the piece size; release of the pinned masters (matters under DS_PIN_MEMORY_BACKEND=native)

Design notes:

• Why the cast moves to the device: pinning alone does not make the copy faster. A blocking copy_ from pinned
  float32 into bf16 still casts on the CPU into a new pageable tensor (136.2 ms against 136.5 ms from pageable memory,
  300M elements, in an earlier measurement on the same setup).

• Why pieces: .to(device) of a whole sub-group makes a float32 device temporary of the sub-group's size,
  3.73 GiB at the default sub_group_size of 1e9. Pieces of 2**26 elements keep it at 0.25 GiB at the same speed.

  Measured on one H200, torch 2.8.0, 12 CPU threads, copying one partition back into its bf16 device partition:

  Copy back	300M elements	1e9 elements	Extra device memory at 1e9
  copy_ from pageable float32 (master)	132.6 ms	483.3 ms	0
  pinned float32, .to(device), then copy_ (fix pagable h2d memcpy #5301)	22.1 ms	73.6 ms	3.73 GiB
  the same in pieces of 2**26 elements (this PR)	22.1 ms	73.7 ms	0.25 GiB
  earlier version of this PR: two pinned 2**26-element bf16 staging buffers	19.4 ms	73.6 ms	0

• Same values: the device cast rounds to nearest even, as the CPU cast does. The new test compares the trained
  parameters bitwise against pin_memory: false.

• Stream order: the pieces go on the current stream. 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
  that stream does not wait for the main stream.

• Host synchronization: one wait per sub-group, at the end of its copy. The old path also 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 column; all runs in one job on the same node.

	master 0f8c9d7f1	earlier version of this PR (staging buffers)	this PR
step time	12.08 s	10.20 s	9.42 s
step time with wall_clock_breakdown	11.96 s	-	9.42 s
rank 0: forward / backward	1.19 / 3.71 s	-	1.19 / 3.67 s
rank 0: step / of which the sub-group loop (optimizer_step)	7.06 / 6.05 s	-	4.56 / 3.52 s
peak device memory (rank 0)	34.66 GB	34.66 GB	34.66 GB
loss at steps 1 / 2 / 8	0.022649 / 0.976290 / 1.746374	0.022649 / 0.976256 / 1.746135	0.022649 / 0.976285 / 1.746306

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

This version is faster than the earlier one in training (9.42 against 10.20 s), although the two copies take the
same time alone at 1e9 elements. A likely reason, not measured: the earlier version cast on the CPU, which the CPU
optimizer of all 8 ranks on the node also uses; this version casts on the device.

Earlier version of this PR (staging buffers) on MiniMax-M3 (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.

These were not repeated with this version.

Tests

Updated: 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 requires every CPU master partition to be pinned with pin_memory: true, and none with pin_memory: false.
• It requires that the pinned run used the new copy. Otherwise the test would compare the direct copy with itself.
• The piece size is set to 1,000 elements, so each sub-group is copied in many pieces.

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 0f8c9d7f1 + only the new test: 2 failed, AttributeError: no attribute '_copy_pinned_master_to_device'
# 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/        # 47 passed, 1 skipped (torch_npu)

The test is in tests/unit/v1, so PR CI selects it: the modal-torch-latest selector (ci/tests_fetcher.py)
picked tests/unit/v1/zero/test_zero.py for the earlier version, which changes the same two files.

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 with the earlier version only
ZeRO-3, pin_memory false	unchanged path (not pinned, direct copy); the new test's reference
bf16	yes. fp16: not run (same path; the device cast also rounds to nearest even)
ZeRO-3, NVMe optimizer offload	not run. Unchanged: NVMe-swapped sub-groups are not pinned here and keep the direct copy
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_memory, is_pinned, Event, unpin_memory)

No config, checkpoint format 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!

[sfc-gh-truwase]

Thanks. I recommend mirroring #5301: when offload_optimizer.pin_memory is enabled, pin the final CPU FP32 master partition at construction—after clone().detach()—and explicitly transfer it to the accelerator before the low-precision copy. This is simpler and consistent with ZeRO-1/2 semantics, though it accepts model-scale pinned memory.

One correction: the linked initialization code pins fp32_partitioned_groups_flat[i].grad, not its .data. The flat FP32 master weight currently remains pageable.

[pengdurice]

Oh, @sfc-gh-truwase thanks for the review, I just redid the code to match #5301. Also thank you for the comment of the master weights are not pinned. this is also implemented in the current version. (also sorry for the mis-understanding of the code base where I thought the master was already pinned)Thank you!

[delock]

Hi @sfc-gh-truwase if you have no further comments I'll help merge this PR, thanks!