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

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):

Configuration	Backend	First Loss	Final Loss	Max Diff vs. Baseline
Pre-PR Baseline	Original inline Adam	11.936906	11.518928	0
PR #8278	FusedAdam	11.936906	11.518935	1.14e-5
PR #8278	PyTorch AdamW	11.936906	11.518935	1.05e-5
PR #8278	CPUAdam offload	11.936906	11.461168	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 #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 (#8278).
• Implement 2D unflattening and bounded all-gather buffers for ZeRO-1/2 (stage_1_and_2.py).
• Implement ZeRO-3 CPU offload Muon support with proper state/gradient synchronization (stage3.py).
• (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.
• Support Per-Head Muon with de-interleaving / re-interleaving for fused & interleaved QKV layouts (leveraging CPU L2 cache locality).

---

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

Start by reading MuonWithAuxAdam and the CPU offload paths in stage_1_and_2.py and stage3.py, then review the numerical discrepancy described for CPUAdam. The work is complete when 2D unflattening, bounded all-gather buffers, ZeRO-3 synchronization, and the listed Muon layout support are implemented and unit tests confirm loss curves match the non-offloaded ZeRO baseline.