Fix Muon optimizer under ZeRO CPU offload and bound gather buffers - #8464

Draft

jinyouzhi wants to merge 13 commits into

deepspeedai:masterfrom

jinyouzhi:muon-cpu-offload-fix

Commits on Aug 27, 2026

Fix Muon CPU optimizer offload
jinyouzhi
committed
8103ab3
Improve the comm overhead, bound Muon offload gather buffers
jinyouzhi
committed
6897349

Commits on Sep 8, 2026

Fix ZeRO-3 offload device mismatch, ZeRO-1/2 gather nesting, and cache eviction

jinyouzhi
and
Copilot
committed
b139b8f
Merge remote-tracking branch 'origin/master' into muon-cpu-offload-fix
jinyouzhi
committed
93c5d7a
Keep Muon momentum buffer resident when save_muon_momentum_buffer_in_memory is set

jinyouzhi
and
Copilot
committed
ca25ef9
fix(zero): preserve Muon update scaling under CPU offload with loss scaling
jinyouzhi
committed
b4ac49e
Fix Muon CPU offload gradient clipping norm and ZeRO-3 NVMe momentum residency

jinyouzhi
and
Copilot
committed
a12f7b6

Commits on Sep 9, 2026

Fix ZeRO-3 NVMe unswapped gradient fragment loss and support pipelined swapper

jinyouzhi
and
Copilot
committed
fce6c97
Add ZeRO-3 NVMe mixed swapped/unswapped numerical equivalence test

jinyouzhi
and
Copilot
committed
a55e543
Fix ZeRO-3 Muon parameter selection and use independent full-gradient reference

jinyouzhi
and
Copilot
committed
12fa5de
Add 2-rank ZeRO-3 NVMe Muon test with DP gradient averaging oracle

jinyouzhi
and
Copilot
committed
9f74ac9
Keep ZeRO-3 Muon momentum in fp32 and tighten the NVMe oracles

jinyouzhi
and
Copilot
committed
6cf979b
Restore ZeRO-3 resident Muon momentum when loading a checkpoint

jinyouzhi
and
Copilot
committed
5a596a5