[RFC] Enable Correct and High-Performance Muon Optimizer under ZeRO CPU Offload (Stage 1/2/3) #8463

1. Motivation & Background

In PR #8278, we extended the auxiliary optimizer in MuonWithAuxAdam to support optimized backends such as FusedAdam and CPUAdam. During our end-to-end convergence validation (Qwen-style 318M model, 100 steps, Alpaca dataset), we observed that while FusedAdam and PyTorch AdamW achieved near-identical loss trajectories compared to the baseline (Max Diff ~ 1e-5), CPU offload exhibited a significant loss divergence (Max Diff ~ 9.37e-2):

Special thanks to @delock for the valuable insights and prior exploratory work in PR #7939, which laid important groundwork for understanding Muon under ZeRO offload.

2. Root Cause Analysis

Muon relies on 5~6 iterations of Newton–Schulz polar decomposition over $\ge 2\text{D}$ weight matrices ($| \Delta W |_2 = 1$).

Under DeepSpeed ZeRO CPU Offload:

1. Geometric Structure Loss: ZeRO flattens model parameters and gradients into 1D partitioned chunks (flat_partition) before copying D2H. Muon’s matrix-level update requires full 2D gradients; flattening and slicing breaks the matrix geometric properties.
2. Silent Fallback: When offloaded, the parameter update path effectively degenerates to 1D scalar/element-wise updates, completely bypassing Newton–Schulz orthogonalization and first-order momentum tracking.

3. Review of Past Attempt: PR #7939

In PR #7939 (feat(zero2): add CPU offload support for Muon optimizer), a CPU offload path for ZeRO-1/2 was proposed:

• Mechanism: The momentum buffer was hosted in CPU memory. During step(), momentum and gradients were transferred back from CPU to GPU (H2D), Newton–Schulz orthogonalization ran on the GPU, and the updated momentum was copied back to CPU (D2H).
• Why it was closed: As noted by @delock in PR #7939 comment, the frequent ping-pong PCIe communication overhead (D2H(grad) $\to$ H2D(momentum) $\to$ GPU NS $\to$ D2H(new momentum) $\to$ H2D(weights)) severely degraded training throughput, rendering ZeRO-2 CPU offload impractical. Moreover, ZeRO Stage 3 was left unsupported.

4. Proposed Solution & Architecture

To address both numerical correctness and PCIe communication bottlenecks, we propose an end-to-end design covering ZeRO Stage 1, 2, and 3:

A. Metadata Propagation & Shape Reconstruction

• Topology-Aware Slicing: Preserve and pass parameter geometric metadata (2D/3D shapes and head dimensions) into the CPU offload path.
• Separation of Concerns:
  • 1D parameters (Bias, LayerNorm, Embeddings) follow the standard 1D flat partition path via CPUAdam.
  • 2D weight matrices are reconstructed into full logical matrices before executing Newton–Schulz, and sliced back to rank partitions afterward.

B. Bounded All-Gather Buffers & Comm Overhead Reduction

• For multi-GPU ZeRO-1/2/3, unflattening requires gathering rank partitions. We introduce an LRU-managed _muon_allgather_buffers cache (capped at 256MB) using OrderedDict to eliminate memory fragmentation and avoid unbounded VRAM allocation spikes.

C. Local CPU Execution & Hardware Acceleration Exploration

• Zero GPU Optimizer Overhead: GPU VRAM strictly maintains only the FP16/BF16 weights and forward activations. FP32 master weights, momentum buffers, and optimizer states reside exclusively in CPU host memory.
• Minimal PCIe Traffic (Single-Round Transfer):
  1. Forward / Backward runs entirely on GPU.
  2. Single D2H transfer of gradients to CPU host memory.
  3. Newton–Schulz orthogonalization (5~6 GEMMs) is executed locally on CPU.
  4. Single H2D transfer of updated FP16/BF16 weights back to GPU.
• Potential Acceleration (AMX Concept): Since multiple GEMMs on CPU can be compute-intensive under plain AVX-512, we could explore leveraging hardware matrix acceleration on modern server CPUs (such as Intel Xeon AMX / TMUL for BF16) to speed up host-side Newton–Schulz iterations, which may help mitigate host compute latency without adding PCIe ping-pong.

D. Per-Head Muon (MuonSplit) & Interleaved Layout Support (CPU Cache Locality)

• Problem with standard Muon & GPU Strided Slice: Standard Muon couples all heads in large Q/K/V projections ($[H \times D, I]$). However, in fused & interleaved layouts ($[Q_0, K_0, V_0, Q_1, K_1, V_1, \dots]$), slicing and permuting non-contiguous head chunks on GPU incurs significant kernel launch overhead and memory fragmentation (as noted in PR [muon] Per-head Newton-Schulz for attention projections #8384).
• CPU Host Advantage: Under CPU offload, master weights and optimizer states reside in host RAM. Single-head chunks (typically tens of KB to 1MB) fit squarely within modern CPU L2/L3 caches (e.g., 23MB L2 per core on modern Xeon). Strided de-interleaving, per-head Newton–Schulz iterations, and re-interleaving can be executed in-place within host cache without GPU kernel launch overhead or PCIe ping-pong, transferring only contiguous tensors over PCIe.

5. Roadmap

• Investigate numerical discrepancy under CPUAdam (Enable optimized Adam backend for MuonWithAuxAdam optimizer #8278).
• Implement 2D unflattening and bounded all-gather buffers for ZeRO-1/2 (stage_1_and_2.py) (Fix Muon optimizer under ZeRO CPU offload and bound gather buffers #8464; follow-up fix Stage Muon's momentum once per partition in ZeRO-1/2 #8633).
• Implement ZeRO-3 CPU offload Muon support with proper state/gradient synchronization (stage3.py) (Fix Muon optimizer under ZeRO CPU offload and bound gather buffers #8464; follow-ups Fix ZeRO-3 Muon all-gather for mixed parameter shapes #8628, Run Muon once per optimizer step under ZeRO-3, not once per micro-batch #8600).
• Cut the PCIe cost of the device-side Newton–Schulz offload path ([Muon] CPU offload: pin the momentum and overlap the host writeback #8858): pin the Muon momentum under offload_optimizer.pin_memory in ZeRO-1/2/3 (including after a checkpoint load, which replaces it with pageable tensors), and copy each update and momentum back to the host on a side stream while the next one is gathered and orthogonalized. Results are bitwise unchanged; measurements are in the PR.
• (Follow-up) One-step-stale Muon update, to hide orthogonalization behind the next forward/backward. This changes the optimizer (the update applied at step t comes from the gradient of step t-1), so it needs convergence evidence, a checkpointed pending update, and a defined interaction with overflow skipping, gradient clipping and accumulation boundaries. Only worthwhile where Newton–Schulz cannot run on the device.
• (Follow-up) Newton–Schulz with per-iteration optimized coefficients (Polar Express, arXiv:2505.16932) as an ns_method option, so fewer iterations reach the same polar-factor accuracy. This lowers orthogonalization cost on any device and, together with per-head Muon, is the main lever if host-side Newton–Schulz is revisited. Its results differ from the current quintic iteration (as gram already does), so it needs accuracy and loss-curve validation against the default.
• (Exploratory) Profile host-side Newton–Schulz compute and explore CPU GEMM / AMX acceleration feasibility.
• Add unit tests and verify 100-step loss curves match non-offloaded ZeRO baseline. (Unit tests landed in Fix Muon optimizer under ZeRO CPU offload and bound gather buffers #8464; loss curves for ZeRO-1/2/3: [RFC] Enable Correct and High-Performance Muon Optimizer under ZeRO CPU Offload (Stage 1/2/3) #8463 (comment))
• Support Per-Head Muon with de-interleaving / re-interleaving for fused & interleaved QKV layouts (leveraging CPU L2 cache locality). (The device-side per-head Newton–Schulz landed in [muon] Per-head Newton-Schulz for attention projections #8384 and [muon] Per-head Muon for linear-attention layers: read head geometry from the owning module #8436, and fused QKV is in [muon] Per-head Newton-Schulz for fused QKV projections #8580; the host-side variant remains, together with the host-side Newton–Schulz item above.)

cc @delock thanks for insights and happy to hear any advices.