[Muon] CPU offload: pin the momentum and overlap the host writeback

[jinyouzhi]

Speeds up Muon under ZeRO CPU optimizer offload (ZeRO-1/2/3) without changing results: the momentum now honors offload_optimizer.pin_memory, and the host writeback of each update overlaps the next orthogonalization. Part of the roadmap in #8463.

Motivation

With CPU offload, Newton–Schulz runs on the accelerator, so every step copies each Muon gradient and momentum to the device and the update and momentum back. Every other buffer that crosses PCIe on this path is pinned when pin_memory is set; the Muon momentum was not:

• ZeRO-1/2 allocated it with torch.zeros_like, ZeRO-3 with torch.zeros (unpinned_fp32_buffer_momentum).
• Loading a checkpoint replaces optimizer state with freshly allocated pageable tensors, so even a pinned momentum would lose the pin on resume.

This is the same per-step data flow as #7939, which added Muon CPU offload to ZeRO-1/2 and was closed because the offload path was too slow (cc @delock). That design is now on master; this PR targets its two costs, the pageable momentum copies and the serialized writeback.

Pageable copies ran at ~10 GB/s here versus ~43 GB/s H2D / ~57 GB/s D2H pinned. The writeback was also strictly serialized with the next batch's gather and Newton–Schulz.

Changes

• Pin the momentum (stage_1_and_2.py, stage3.py): allocate it with pin_empty(_like) when offload_optimizer.pin_memory is set, and re-pin it at the point of use if it is not pinned (after a checkpoint load). ZeRO-3 momentum that the NVMe swapper manages is left to the swapper; the save_muon_momentum_buffer_in_memory copy is pinned.
• Overlap the writeback (zero/utils.py: OverlappedHostWriteback): the update and momentum of batch i (ZeRO-1/2) or parameter i (ZeRO-3) are copied to the host on a side stream after batch i+1 has been queued, so the D2H overlaps the next gather and Newton–Schulz. Dependencies are enforced with host-side Event.synchronize() rather than a cross-stream wait, since on XPU a stream waiting on another stream's event serialized the two queues (0 overlap in a trace). Sources stay referenced until their copy completes (at most one batch computed and one in flight), and the copies are drained before the host reads them (end of the update, and before an NVMe swap-out). Synchronized devices (CPU accelerator) copy inline as before.

The update math is untouched; results are bitwise identical to master.

Tradeoffs

• Pinned host memory grows by the momentum size (4 B per Muon parameter per rank; 14.5 GB for Qwen3-4B on one rank), only when the user opted into pin_memory.
• Device memory: one more batch of local update/momentum slices can be alive while its copy is in flight (bounded by the existing Muon gather budget). Peak memory is reported below.

Validation

Hardware: 8× Intel Arc Pro B70 (XPU), 2× Xeon Platinum 8457C, 1 TB RAM. Software: torch 2.14.0.dev20260804+xpu, Level Zero / compute runtime 26.27.39122.11.

Tests

New TestMuonOffloadPinnedMomentum in tests/unit/v1/ops/muon/test_muon_checkpoint.py (ZeRO 1/2/3 × pin_memory True/False, world size 2, bf16): after training, the momentum is pinned exactly when pin_memory is set; after save_checkpoint → load_checkpoint → training it is pinned again, and the resumed run matches the uninterrupted one.

cd tests && pytest unit/v1/ops/muon/test_muon_checkpoint.py -k Pinned

Code	Result
master	3 failed (pin_memory=True, stages 1/2/3), 3 passed
this PR	6 passed

Full Muon suite (cd tests && pytest unit/v1/ops/muon/): 250 passed, 10 skipped, 4 failed. The same 4 fail on master on this machine and are unrelated to offload: test_muon_overflow.py::TestMuonMatchesZeroStage0::test_every_matrix_keeps_its_momentum[1,2] (9.7e-4 vs the 1e-5 bound) and test_per_head_muon_accelerator.py::test_per_head_on_a_shard_is_the_shard_of_per_head[2-gram,2-standard] (bitwise equality of sharded vs whole Newton–Schulz on XPU).

Performance

Qwen3-4B shapes (randomly initialized, 36 layers, 4.02B parameters of which 3.63B Muon), seq 2048, micro-batch 1, bf16, gradient checkpointing, gradient clipping 1.0, offload_optimizer: {device: cpu, pin_memory: true}, 1 GPU, one socket (48 cores, numactl -C 48-95 -m 1). Per run: mean of steps 2–6; 3 runs per cell. "Muon apply" is _apply_muon_updates_cpu_offload (gather, Newton–Schulz, writeback), timed with device synchronization.

ZeRO	Code	Iteration (s), runs	Iteration mean	Muon apply mean	Peak device memory
2	master	8.51 / 8.52 / 8.60	8.54	5.13 s	12.09 GiB
2	this PR	5.57 / 5.77 / 5.70	5.68 (−33%)	2.35 s (−54%)	12.09 GiB
3	master	7.66 / 7.65 / 7.64	7.65	3.79 s	19.88 GiB
3	this PR	6.27 / 6.30 / 6.27	6.28 (−18%)	2.22 s (−41%)	20.11 GiB (+0.23)

Breakdown of the two changes (same setup, this PR's code with one change disabled; 3 runs per cell). Pin only forces the writeback synchronous; overlap only leaves the Muon momentum pageable while every other offload buffer stays pinned. The momentum pin status was checked in every run.

ZeRO	Variant	Iteration (s), runs	Iteration mean	Muon apply mean	Peak device memory
2	master		8.54	5.13 s	12.09 GiB
2	pin only	5.92 / 6.01 / 5.88	5.94 (−30%)	2.82 s (−45%)	12.09 GiB
2	overlap only	7.09 / 7.08 / 7.19	7.12 (−17%)	3.85 s (−25%)	12.09 GiB
2	both (this PR)		5.68 (−33%)	2.35 s (−54%)	12.09 GiB
3	master		7.65	3.79 s	19.88 GiB
3	pin only	6.32 / 6.22 / 6.30	6.28 (−18%)	2.57 s (−32%)	19.88 GiB
3	overlap only	6.57 / 6.50 / 6.61	6.56 (−14%)	2.91 s (−23%)	20.11 GiB
3	both (this PR)		6.28 (−18%)	2.22 s (−41%)	20.11 GiB

Pinning is the larger effect. The overlap adds a further 0.47 s (ZeRO-2) and 0.35 s (ZeRO-3) of Muon apply on top of pinning. Under ZeRO-2 that shows up in the iteration time (5.94 → 5.68 s). Under ZeRO-3 the iteration mean does not move (6.28 s in both); the gain is within this setup's iteration noise. The +0.23 GiB ZeRO-3 peak comes from the overlap (the in-flight parameter), not from pinning.

Why ZeRO-3 peaks higher than ZeRO-2 here (pre-existing, same on master). Peak device memory per phase, this PR, single GPU:

Phase	ZeRO-2	ZeRO-3	ZeRO-3, sub_group_size: 2.5e8
forward	11.51 GiB	11.72 GiB	11.72 GiB
backward	12.09 GiB	12.30 GiB	12.30 GiB
step	9.05 GiB	20.11 GiB	11.86 GiB

With one rank the ZeRO-3 partition is the whole model, so forward and backward match ZeRO-2. The difference is in the step: ZeRO-3's Muon offload apply copies the gradients and momentum of a whole sub-group to the device and gathers them at once (default sub_group_size 1e9 elements), while ZeRO-1/2 processes Muon parameters in batches bounded by a gather budget. With sub_group_size: 2.5e8 the step peak drops to 11.86 GiB, and the iteration time is unchanged (6.10 s). Batching the ZeRO-3 apply the same way as ZeRO-1/2 is left out of this PR; with more ranks the per-rank partitions shrink but the gathered full gradients and momentum of a sub-group do not.

Correctness of the overlap: in the same process, the Muon apply was run twice on identical host state, once with the writeback forced synchronous and once overlapped, and the resulting gradients and momentum were compared with torch.equal (ZeRO-2 and ZeRO-3, 4 layers). They are bitwise equal on every step after the first. On the first step two synchronous runs already differ from each other by the same amount (max 4.9e-4 in one gradient tensor), so that is first-call nondeterminism of the existing path, not the overlap. Losses are not bitwise reproducible across separate processes on master either (they differ from step 1 on), so cross-run loss equality is not used as evidence; the trajectories of master and this PR stay within master's own run-to-run spread.

Reproduction

The benchmark builds Qwen3ForCausalLM from transformers with the shapes above, wraps _apply_muon_updates_cpu_offload with a synchronized timer, and runs 7 steps:

RANK=0 LOCAL_RANK=0 WORLD_SIZE=1 MASTER_ADDR=127.0.0.1 MASTER_PORT=29677 \
MALLOC_MMAP_THRESHOLD_=17179869184 MALLOC_TRIM_THRESHOLD_=17179869184 MALLOC_TOP_PAD_=1073741824 \
STAGE={2,3} numactl -C 48-95 -m 1 python bench_pin.py

Not validated: NVMe optimizer offload with Muon, and multi-node runs.