💻 computer science

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

MatrixFSDP enables communication-free large-scale training with matrix optimizers like Muon under ZeRO-3 sharding by reorganizing parameter placement so that each 2D weight matrix resides entirely on a single rank, thereby eliminating the need for costly matrix reconstruction during optimizer steps while maintaining memory efficiency and achieving significant latency reductions.

Original authors: Ming Gao, Yanwu Xu, Hao Zhang

Published 2026-07-08✓ Author reviewed ⓘ

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ming Gao, Yanwu Xu, Hao Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a massive team of robots (a computer cluster) how to write a novel. The robots are working together to learn from a huge book of text. To do this efficiently, they use a special "learning rule" called Muon.

The Problem: The "Whole Picture" vs. The "Puzzle Pieces"

Normally, when robots learn, they break the massive book into tiny puzzle pieces. Each robot holds just a few pages (a "shard"). This is great for saving memory because no single robot has to carry the whole book. This method is called ZeRO-3.

However, the Muon learning rule is a bit picky. It doesn't want to learn from just a few pages at a time. To do its job perfectly, it needs to see the entire 2D page (the whole matrix) at once to understand the relationships between words.

The Conflict:

ZeRO-3 says: "We only have puzzle pieces."
Muon says: "I need the whole page to learn."

The Old Solutions (The Bad Options):

The "Reconstruction" Method: Every time the robots need to learn, they stop, gather all the puzzle pieces from every robot, glue them back together to make the full page, Muon learns, and then they immediately tear the page apart again.
- The Downside: This is like a group of people constantly stopping their work to assemble a giant jigsaw puzzle, just to take it apart again. It wastes a huge amount of time and energy (communication) every single step.
The "Full Copy" Method: Instead of sharing pieces, every robot keeps a full copy of the entire book. Muon can learn instantly because everyone has the whole picture.
- The Downside: This requires so much memory that if the book gets too big, the robots' brains (GPUs) explode. They run out of space.

The New Solution: MatrixFSDP

The authors of this paper, MatrixFSDP, found a clever third way. They didn't change the learning rule (Muon) or force everyone to carry the whole book. Instead, they changed who holds the book.

The Analogy: The "Specialized Librarian"

Imagine a library where the books are usually chopped up and distributed among all the librarians.

MatrixFSDP's Idea: For every single "page" (matrix) of the book, they appoint one specific librarian to be the "Owner."
- This Owner holds the entire, complete page.
- All the other librarians hold nothing (an empty space) for that specific page.
- For the parts of the book that don't need the special Muon rule, they stick to the old "puzzle piece" method.

How it Works in Practice:

During Learning (The Optimizer Step): Since the "Owner" already has the whole page, Muon can learn immediately. No one has to gather pieces or glue anything together. It's like the librarian just reading the book right off their desk. Zero communication is needed.
During Reading/Writing (Forward/Backward Passes): When the robots need to read or write to the book, they temporarily bring the pieces back together, do their work, and then immediately put them back into their "Owner" slots.

Why This is a Big Deal

The paper claims this approach solves the biggest bottleneck in training large AI models:

Speed: Because they stopped the constant "gluing and tearing" of pages, the learning step became incredibly fast. On a single computer node, it was 4.2 times faster. On a large cluster of 8 nodes, it was 54.6 times faster because the old method was wasting time sending data across the network between computers, while MatrixFSDP keeps the data local.
Memory: Unlike the "Full Copy" method, MatrixFSDP still only uses the memory of the "puzzle piece" method. It allows them to train models that are too big for the "Full Copy" method to handle.
Accuracy: They proved that because the "Owner" gets the exact same data as if they had gathered the whole book, the learning results are identical to the perfect, slow method.

Summary

MatrixFSDP is like reorganizing a team so that the person who needs the whole document to do their job is the one who actually holds the whole document. Everyone else holds nothing for that specific task. This eliminates the need to constantly pass documents back and forth, making the team work much faster without needing bigger desks (more memory) to hold everything.

Technical Summary: MatrixFSDP

Problem Statement

Large-scale language model training increasingly relies on matrix-structured optimizers like Muon, Shampoo, and SOAP. Unlike coordinate-wise optimizers (e.g., AdamW), these methods require the complete 2D gradient matrix as input to perform operations such as Newton–Schulz orthogonalization or matrix preconditioning. This creates a fundamental systems mismatch with FSDP (Fully Sharded Data Parallel) and ZeRO-3, which shard parameters, gradients, and optimizer states across ranks to maximize memory efficiency.

Under standard ZeRO-3, each rank holds only a slice of the model. To apply a matrix optimizer, the system must reconstruct the full 2D gradient matrix at every optimizer step. This necessitates a weight-sized all-gather communication after the backward pass. This communication is a critical path bottleneck that cannot be overlapped with layer computation, leading to significant latency that scales poorly with node count and model size.

Existing solutions are incomplete:

Reconstruction (Stock FSDP2-Muon): Maintains ZeRO-3 memory savings but incurs heavy communication overhead at every optimizer step.
ZeRO-1 Owner Placement: Assigns whole matrices to specific ranks (owners) to make the optimizer step local, but requires full parameters to be resident on every rank, abandoning ZeRO-3 memory efficiency.

The goal is to achieve whole-matrix optimizer inputs without sacrificing ZeRO-3-scale parameter memory.

Methodology: MatrixFSDP

MatrixFSDP introduces a third path: Owner-shaped ZeRO-3 shards. Instead of changing the optimizer algorithm or abandoning sharding, it redefines where shards reside within the data-parallel group.

1. Owner-Shaped Placement

For every 2D weight matrix in the model, MatrixFSDP designates a single data-parallel rank as the owner.

Owner Rank: Holds the complete 2D matrix.
Non-Owner Ranks: Hold an empty shard for that specific parameter.
Non-Matrix Tensors: 1D vectors and other non-2D tensors are packed into a "tail" and assigned to a separate owner, where they continue to use the standard AdamW optimizer.

This ensures that after the backward pass and gradient reduction, the full Muon input (the complete 2D gradient) naturally resides on the owner rank. Consequently, the Newton–Schulz update runs locally on the owner with zero optimizer-step matrix communication.

2. Runtime Challenges and Solutions

This uneven layout (some ranks holding full matrices, others holding empty shards) violates the equal-size collective assumptions of standard FSDP2. MatrixFSDP addresses this through four key runtime mechanisms:

MatrixShard Metadata & Global Planner: A global planner assigns owners to balance resident bytes and optimizer workload across the cluster. This metadata tracks which rank owns which tensor role, ensuring the system remains a single coherent ZeRO-3 layout rather than a collection of special cases.
Owner-Segment P2P Collectives: Standard all-gathers are replaced with deterministic, owner-segment collectives. During forward/backward materialization, only non-empty segments are transferred from the owner to ranks needing them. During gradient reduction, results are written only to the non-empty owner segments. This avoids padding empty ranks and eliminates full-matrix all-gathers.
Owner-Buffer Pinning: To handle the FSDP2 lifecycle (unshard $\to$ compute $\to$ reshard), MatrixFSDP introduces a "pinned" state. When a buffer is shrunk after the forward pass, the storage is not returned to the shared pool but pinned. This allows the backward pass to "revive" the full matrix in the same memory location, preventing autograd view errors and ensuring the owner can perform the update before releasing the buffer.
Checkpointing and Routing: Checkpointing (DCP) and optimizer state routing use the MatrixShard metadata. Optimizer states are routed to the correct owners (Muon for matrices, AdamW for tails), and checkpoints save the specific placement layout to support resharding across different world sizes.

Key Contributions

Owner-Shaped ZeRO-3: A novel placement strategy where 2D matrices are whole on a single owner and empty elsewhere, enabling local matrix optimization within a sharded memory regime.
Communication-Free Optimizer Steps: By aligning the gradient reduction destination with the optimizer's input requirement, MatrixFSDP eliminates the weight-sized all-gather bottleneck previously inherent to matrix optimizers in ZeRO-3.
Runtime Infrastructure: A suite of runtime components (planner, custom collectives, buffer pinning) that make uneven owner layouts efficient and correct within the FSDP2 framework.
Generality: The approach is optimizer-agnostic, demonstrated to work with Muon, Shampoo, and SOAP without requiring distributed variants of the preconditioners.

Experimental Results

Evaluated on 64 NVIDIA A100 GPUs (8 nodes) with various model sizes (up to 32B parameters):

Optimizer-Step Latency: MatrixFSDP reduces optimizer-step latency by 4.2× on a single node and 54.6× on eight nodes compared to stock FSDP2-Muon. The speedup grows with scale because the baseline's communication crosses the inter-node fabric, while MatrixFSDP's step remains local.
End-to-End Speedup: Achieves up to 2.15× end-to-end training speedup.
Memory Efficiency: MatrixFSDP maintains ZeRO-3 memory levels. For a 32B model, it uses 10 GB/rank, whereas a ZeRO-1 owner placement approach exceeds the 80 GB GPU limit (projected >54 GB/rank at 14B). It is 3–7× more memory-efficient than "gather-once" FSDP2-Muon approaches.
Correctness: The training trajectory matches a full-matrix DDP Muon reference with high precision (printed loss difference $\leq 0$ ), confirming that the owner receives the exact gradient matrix required for the update.

Significance and Claims

The paper claims that MatrixFSDP solves the "systems mismatch" between matrix optimizers and parameter sharding. Its primary significance lies in enabling large-scale training with matrix optimizers (which offer better token efficiency and convergence) without incurring the prohibitive communication costs or memory overheads of previous approaches.

The authors emphasize that MatrixFSDP is particularly effective for multi-node ZeRO-3 runs with large shard spans, where the cost of reconstructing matrices at every step would otherwise dominate training time. The paper modestly notes that the gains are smaller on single-node runs or under fixed-model strong scaling where the number of owner roles becomes sparse, potentially making the owner fanout a new bottleneck. The work does not claim to replace coordinate-wise optimizers for non-matrix parameters but rather provides a unified, efficient path for handling the 2D matrix components of modern models.