Search papers, labs, and topics across Lattice.
This paper introduces MatrixFSDP, a novel approach to matrix optimization that addresses the limitations of existing systems by allowing for communication-free updates in large-scale training scenarios. By reconfiguring the placement of ZeRO-3 shards, MatrixFSDP enables the Muon optimizer to operate on complete matrices without incurring the heavy communication costs typically associated with matrix reconstruction. The results show that MatrixFSDP achieves a 4.2x reduction in optimizer-step latency on a single node and a remarkable 54.6x reduction on eight nodes, while also enabling training of larger models that exceed the memory constraints of traditional methods.
MatrixFSDP slashes optimizer-step latency by over 54x while enabling training on models that exceed 80 GB GPU limits.
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.
