MatrixFSDP: communication-free matrix optimizers under ZeRO-3 parameter sharding

2026-07-07

AI Summary

This work addresses the incompatibility between ZeRO-3’s parameter sharding mechanism and matrix-based optimizers like Muon, which require full weight matrices and thus incur excessive communication or memory overhead. To resolve this, the authors propose MatrixFSDP, a framework that restructures tensor sharding: matrix tensors are fully owned by a single data-parallel rank while other ranks hold empty shards, and non-matrix tensors are packed into a tail rank for standard AdamW optimization. This design enables zero-communication matrix optimization without altering optimizer logic. Key innovations include MatrixShard metadata, a load-aware Owner planner, deterministic point-to-point communication, and checkpoint resharding. Experiments on 64×A100 GPUs show that, compared to FSDP2-Muon, MatrixFSDP reduces optimizer step latency by 4.2× on a single node and 54.6× across eight nodes, achieving a 2.15× end-to-end training speedup and enabling large-model training beyond the 80GB GPU memory limit.

Abstract

Matrix optimizers such as Muon are attractive for large-scale training because they can improve convergence and token efficiency over coordinate-wise optimizers. Muon does this by orthogonalizing momentum-smoothed matrix updates with Newton-Schulz, producing spectrum-balanced updates that require the complete 2D matrix as input. This exposes a systems mismatch: FSDP/ZeRO-3 saves memory by making the optimizer see shards, not whole matrices. Existing systems therefore either reconstruct matrices at every optimizer step, paying weight-sized communication after backward, or make the update local by using ZeRO-1 owner placement with full parameters resident. MatrixFSDP takes a third path: it changes where ZeRO-3 shards live, not the optimizer being computed. For each 2D weight, one data-parallel rank owns the whole matrix and the other ranks hold empty shards; non-matrix tensors are packed into tail owners and stay on AdamW. The ordinary backward reduction then lands the full Muon input on the owner, so Newton-Schulz runs locally with no optimizer-step matrix collective. Forward and backward still materialize and reshard parameters; the runtime challenge is to make that uneven layout efficient and correct. MatrixFSDP does so with MatrixShard metadata, a balance-aware owner planner, deterministic owner-segment P2P collectives, owner-buffer pinning, and owner-shard checkpoint resharding. The resulting update matches full-matrix Muon while preserving ZeRO-3-scale memory: on 64 A100s, MatrixFSDP reduces optimizer-step latency over stock FSDP2-Muon by 4.2x on one node and 54.6x on eight nodes, reaches up to 2.15x end-to-end speedup, and runs model sizes where ZeRO-1 owner placement exceeds an 80 GB GPU.

Problem

• matrix optimizers
• ZeRO-3
• parameter sharding
• communication overhead
• large-scale training

Innovation

• MatrixFSDP
• ZeRO-3
• matrix optimizer
• communication-free
• parameter sharding