MatrixFSDP: communication-free matrix optimizers under ZeRO-3 parameter sharding
July 07, 2026
1 Introduction↩︎
As language models grow, training systems must do more than fit larger parameter counts: they must also reduce the time and tokens needed to reach a target quality. Matrix-structured optimizers are one promising direction. Muon [1], for example, improves token efficiency over AdamW and has been adopted in recent large-scale training runs [2], [3]; Shampoo [4], [5] and SOAP [6] pursue the same goal with matrix preconditioners. The catch is that these optimizers ask the distributed runtime for a different object. FSDP2 / ZeRO-3 [7], [8] shard parameters, gradients, and optimizer state into per-rank slices, which is exactly the right contract for coordinate-wise AdamW. Muon instead applies Newton–Schulz to a whole 2D gradient: a row-band is not the same optimizer input as the full matrix. Step efficiency thus becomes a ZeRO-3 systems problem.
Existing systems expose two incomplete choices. FSDP2-Muon keeps ZeRO-3 memory but reconstructs each matrix at optimizer time, placing weight-sized communication after backward. ZeRO-1 owner placement makes the matrix operation local, but keeps full parameters resident on every rank. The missing combination is whole-matrix optimizer inputs without abandoning ZeRO-3 parameter memory.
The hard part is the runtime contract around that placement. At the optimizer boundary, Muon should see the same complete 2D gradient as a gathered reference; between steps, no rank should retain a full model copy; and forward/backward must still use FSDP2’s transient materialize/reshard buffers. We therefore present MatrixFSDP, an FSDP2 runtime for owner-shaped ZeRO-3 shards:
\(\triangleright\) Owner-shaped placement. Each data-parallel-sharded 2D weight is whole on one owner and empty elsewhere, while the non-matrix tail is packed as a whole AdamW owner role. Backward reduction deposits the full Muon input on the owner, so the optimizer step needs no matrix gather or redistribute. Tensor-parallel-owned fragments are outside this scope.
\(\triangleright\) Planner and metadata. A global owner planner balances resident bytes and optimizer work across FSDP blocks, while execution remains block-local for FSDP2 prefetch and overlap.
MatrixShard metadata carries the same owner plan through optimizer routing, owner-buffer pinning, DCP save/load, and cross-world-size resharding.
\(\triangleright\) Owner-segment communication. Owner-shaped shards violate the equal-size collective assumption. MatrixFSDP uses deterministic owner-segment P2P collectives that move only nonempty pieces, avoiding both largest-owner padding and full-matrix all-gathers.
Our evaluation shows that MatrixFSDP preserves the optimizer trajectory while removing optimizer-step matrix communication. It matches a DDP Muon reference on real-data loss, reduces optimizer-step time as node count increases, keeps ZeRO-3-scale parameter memory where ZeRO-1 owner placement runs out of memory, and reuses the same owner path for Shampoo and SOAP.
2 Motivation↩︎
In FSDP2 / ZeRO-3 [7], [8], each rank stores only a slice of each parameter. Full weights exist transiently for forward/backward, then are immediately resharded; after backward, each rank again holds only its local parameter and gradient slices. This is exactly enough for AdamW, whose update is coordinate-wise. Once the gradient slice is reduced, the optimizer step needs no further model-sized collective.
Muon [1] changes that contract. For a 2D weight with momentum-buffered gradient \(G\), Muon applies a Newton–Schulz orthogonalization whose steps are built from products such as \(XX^\top\). Those products couple the rows of \(G\): a row-band is not a valid input to the same update. Shampoo [4], [5] and SOAP [6] have the same shape of problem: their preconditioners are also whole-matrix operators. Thus, under ordinary equal ZeRO-3 sharding, the optimizer must first materialize each 2D gradient as a whole matrix.
Figure 1 contrasts the two existing endpoints with the target design point. The common ZeRO-3 endpoint reconstructs each 2D gradient at optimizer time: gather the matrix, run Newton–Schulz, then scatter the update back to ZeRO-3 shards. If the model has \(N\) 2D parameters of \(s\) bytes over \(W\) data-parallel ranks, this moves roughly \(2Ns\frac{W-1}{W}\) extra bytes per rank per step. More importantly, it moves them after backward has finished. Unlike FSDP’s forward/backward all-gathers, this traffic cannot be overlapped with layer compute; it is pure optimizer critical path. For transformers, where 2D weights dominate parameter mass, this cost is large, recurring, and grows with model size and node count.
The opposite endpoint removes that optimizer-time reconstruction by assigning whole matrices to owners. Distributed Shampoo [9] and Canzona’s data-parallel path [10] assign whole matrices to owners, so the matrix operation itself is local. The catch is that these systems operate in a ZeRO-1-style regime with full parameters resident on every rank. They remove optimizer-step communication by giving up the parameter-memory savings that ZeRO-3 was meant to provide. MatrixFSDP targets the missing point in Figure 1: whole-matrix ownership as a ZeRO-3 placement, not as a replacement for ZeRO-3.
3 Approach: Owner-Shaped ZeRO-3↩︎
MatrixFSDP keeps the FSDP training loop but changes the resident shard shape seen after each reshard. For each 2D parameter \(W_i\), a data-parallel rank \(o_i\) is chosen as its owner. In the resharded state, rank \(o_i\) stores the whole matrix and every other rank stores an empty shard for that parameter. The non-2D tensors in the same FSDP unit are packed into a tail role and assigned whole to a tail owner, where they remain on the AdamW path. Across the data-parallel group there is still only one resident copy of each parameter; the difference from ordinary ZeRO-3 is that the copy is assigned by tensor role rather than sliced equally across all ranks. Matrices already fragmented by tensor parallelism are not owner-placed here; they are skipped by the matrix-owner planner and handled by the surrounding TP runtime.
This placement makes the optimizer input local. Backward still produces the same per-rank gradients as ordinary data parallel training, and ZeRO-3 still reduces them into the current local shard shape. Because the local shard of a 2D matrix is the full matrix on its owner, the reduced Muon input \(G_i\) lands on that owner. The owner applies the Newton–Schulz update to the full matrix and updates its local parameter and Muon state; ranks with empty shards do no optimizer work for that matrix. Tail owners apply AdamW to their packed non-matrix tensors. Thus the optimizer step contains no matrix all-gather, broadcast, or redistribute. The next forward pass materializes parameters through FSDP’s usual unshard/reshard lifecycle, now from owner-shaped resident shards.
The price of this layout is that the resident shards are no longer equal-size or evenly present on all ranks. A rank may own several matrices, a large tail, or nothing for a given FSDP unit. If the runtime handles this with ordinary equal-size
collectives, it either pads empty ranks up to the largest owner shard or reconstructs every matrix on ranks that will not optimize it. MatrixFSDP therefore treats ownership as a first-class placement: a global planner chooses owners to balance resident
bytes and optimizer work, MatrixShard metadata records the nonempty owner segments, and custom P2P owner collectives move only those segments during materialization and gradient reduction. The same placement metadata drives optimizer routing
and checkpointing, so the system remains one ZeRO-3 layout rather than a collection of optimizer-specific special cases.
4 FSDP2 Runtime for Owner Shards↩︎
§3 changes only the resident shard shape, but that change cuts across the runtime assumptions that make FSDP2 efficient. Ordinary FSDP2 assumes that every parameter has a nonempty, roughly equal local shard on every rank; its materialization and gradient-reduction paths can therefore use equal-size collectives, and its optimizer and checkpoint code can infer most behavior from the local tensor shape. Owner placement breaks all three assumptions. For a matrix, one rank owns the full local shard, the other ranks own empty shards, and different ranks may own very different total bytes within the same FSDP unit.
MatrixFSDP makes this layout usable by turning ownership into explicit runtime metadata rather than an optimizer-side convention. The metadata first describes which rank owns each matrix and tail tensor, using a planner whose global assignment balances memory while block-local execution preserves FSDP2 prefetch and overlap (§4.1); custom owner-segment collectives then move only nonempty pieces (§4.2); FSDP2’s buffer lifecycle is tightened so owner-shaped shards can still participate in the usual unshard/reshard schedule (§4.3); and the same metadata routes optimizer state and checkpoint shards (§4.4). The goal is not to replace FSDP2’s training loop, but to make an imbalanced owner layout look like a valid ZeRO-3 placement at each boundary where FSDP2 expects a shard.
4.1 MatrixShard: representing ownership↩︎
A standard Shard(0) placement is an equal split of a tensor interval. MatrixFSDP instead records a rank-indexed placement for each flattened parameter interval. A 2D matrix has exactly one full segment, stored on its matrix owner, and empty
segments elsewhere. The non-2D tail in an FSDP unit is packed into one whole segment on a tail owner. Figure 5 contrasts this state with ordinary FSDP2 after reshard: ordinary FSDP2 gives every rank a fraction of
every tensor, while MatrixFSDP gives each tensor role to one owner rank.
The planner is global, but execution is block-local (Figure 2). Planning sees all FSDP blocks selected by the wrap policy. Each block is converted into roles: one role per 2D matrix and one packed non-2D tail. The greedy load is owner-resident numel, a proxy for resident parameter/gradient bytes and Muon work; the resource scorer expands a completed plan into per-rank parameter, gradient, optimizer-state, communication, and workspace estimates. Balance is therefore decided at model scope, not from one block’s local shape. The same owner-byte balance later bounds the fanout bottleneck in §4.2: if one block assigns too many bytes to one owner, that rank becomes the source for materialization and the sink for gradient reduction.
Algorithm 3 summarizes the owner assignment. RoleGreedy is the memory-conservative default; ScopeGreedy pools roles from all blocks to improve global balance; CostAware accepts the scope plan only when the balance gain does not exceed workspace caps. Each emitted block layout is then checked for one contiguous local segment per parameter, so the hot path remains a standard FSDP2 block schedule.
4.2 Owner-segment communication collectives↩︎
Once shards are owner-shaped, the equal-size collective path is the wrong abstraction. Padding every rank to the largest owner segment wastes memory and bandwidth, while all-gathering a matrix to every rank recreates the optimizer-step reconstruction that owner placement was meant to remove. The logical operations FSDP2 needs are narrower: materialize a full parameter from its current owner segments, and reduce the full gradient back into the same owner segments.
4.2.0.1 Bandwidth model.
Let \(W\) be the shard span: the data-parallel ranks over which a parameter is sharded, materialized, and reduced; if world size exceeds \(W\), the remainder is the HSDP replicate dimension. Let \(P_B\) be the parameter bytes in FSDP block \(B\), and \(O_{B,r}\) the bytes in that block owned by rank \(r\). We compare three per-block quantities: \[\begin{align} \gamma_B &= \frac{\max_r O_{B,r}}{P_B/W},\\ C^{\mathrm{FSDP}}_B &\approx P_B \frac{W-1}{W},\\ C^{\mathrm{owner}}_B &= (W-1)\max_r O_{B,r} = \gamma_B P_B \frac{W-1}{W}. \end{align}\] Thus a block with one large owner can be bandwidth-poor (\(\gamma_B \approx W\)); when the planner spreads matrix and tail owners within the block (\(\gamma_B \approx 1\)), the leading bandwidth term matches ordinary FSDP while the optimizer-step reconstruction is removed.
MatrixFSDP implements these operations as owner-segment collectives. For materialization, ranks with nonempty segments send only those segments to the ranks that need the full parameter for the current FSDP unit; in the common matrix case this is an owner fanout. For gradient reduction, each rank contributes its local full-gradient buffer, but the reduced result is written only to the nonempty owner segment. The fast path uses deterministic native send/recv schedules; owner-broadcast and uneven torch collectives are kept as fallbacks. Cross-rank validation checks that every rank derives the same deterministic schedule and segment sizes before communication begins, turning mismatches into explicit errors rather than silent NCCL timeouts.
4.3 Owner-buffer pinning in FSDP2↩︎
FSDP2 unshards before forward, reshards after forward, unshards again before backward, and reduces gradients after backward. MatrixFSDP preserves that lifecycle, but owner-shaped shards expose a storage-level hazard: autograd saves views into the transient full-parameter buffer. Post-forward reshard may shrink that buffer in place, but it cannot return the storage to the shared pool. Before backward, the unit resizes and refills the same storage, so saved views are revived in place (Figure 4).
The implementation adds a pinned state for shrunk buffers that are pending backward. A central guard releases storage only after the owning unit finishes backward and optimizer-ready local shards have been restored. This lets owner-shaped shards coexist
with FSDP2’s materialize-and-reshard discipline without setStorage ... size 0 errors or rank divergence from reviving a reused buffer.
4.4 Optimizer and checkpoint state↩︎
After backward, every FSDP unit is returned to owner-shaped local shards. Optimizer construction then uses MatrixShard metadata, not tensor shape alone: a Muon parameter is visible to the matrix optimizer only on the rank that owns the full
2D shard; ranks with empty shards skip it; and tail-owner shards are routed to AdamW. Since gradient reduction has already deposited each full 2D gradient on its owner, Newton–Schulz runs without optimizer-step collectives.
Checkpointing follows the same placement. DCP save writes each local shard as tensor payload plus a layout sidecar containing the MatrixShard placement, per-rank segments, and replica metadata. Same-layout load validates the recorded
placement before copying payloads into local shards. Reshard load reconstructs each parameter on a canonical \([0,\text{numel})\) axis from the saved segments, then repacks it into the target owner assignment and world
size. Owner-Muon state and tail-owner AdamW state use the same path, so changing world size or owner plan does not require optimizer-specific checkpoint logic.
5 Evaluation↩︎
We evaluate on a multi-node cluster (Sec. 5.1). The three headline results: the optimizer step is communication-free and its speedup over the strongest FSDP2-Muon baseline grows with node count (Sec. 5.5); correctness matches DDP Muon (fp32 tolerance; bf16 printed loss) (Sec. 5.3); and MatrixFSDP holds ZeRO-3 per-rank memory, far below the ZeRO-1 owner placement baseline (Sec. 5.4).
5.1 Setup↩︎
Hardware/software. All GPU experiments run on 8 nodes \(\times\) 8 NVIDIA A100-SXM4-80GB GPUs (64 GPUs) with NCCL over InfiniBand, using PyTorch FSDP2 fully_shard, DTensor, and
torch.optim.Muon. Models/data. Latency and memory use bf16 Qwen-style decoder-transformer shapes (hidden 4096, intermediate 16384, 32 heads, sequence 4096) with activation checkpointing; the weak-scaling sweep uses
16/32/64/128 layers for shard spans 8/16/32/64; here shard span equals world size, so the HSDP replicate dimension is one. Convergence uses a 12-layer decoder-style LM (hidden 768, intermediate 3072, 12 heads, sequence 1024) on a flat uint16 WikiText
stream with GPT-2 vocab 50,257. The embedding and untied LM head are 2D Muon weights; LayerNorm and other non-2D tensors use AdamW. Each convergence step is batch 8 per rank on 64 ranks (524,288 tokens/step), for 10,000 steps (5.24B tokens), fixed seed 0
and rank-local data streams. Optimizers/timing. Muon uses lr \(10^{-3}\), weight decay \(10^{-2}\), momentum 0.5, 2 NS steps, and match_rms_adamw; AdamW uses lr
\(10^{-3}\), weight decay \(10^{-2}\), betas (0.9, 0.999), eps \(10^{-8}\), with no schedule. Latency uses 3 warmup and 10 measured steps; we synchronize
after zero_grad, forward, backward, and optimizer step, then report max-rank phase time. Memory is max-rank peak CUDA allocation after warmup (not reserved cache): it includes bf16 parameter shards and full-parameter buffers, gradients,
optimizer state, MatrixFSDP communication workspaces, and activations from the fixed checkpointed run; none of the modes uses fp32 master weights. Baselines: stock FSDP2-Muon uses PyTorch torch.optim.Muon over
FSDP2-sharded DTensors; gather-once FSDP2-Muon gathers each matrix gradient once, then runs local NS; MatrixFSDP = owner-Muon.
5.2 Optimizer-step and end-to-end speedup↩︎
Table 1 reports one-node latency for a 16-layer, 3.2B-parameter transformer from the weak-scaling sweep (shard span 8). MatrixFSDP’s optimizer step is communication-free—\(87\) ms, versus \(367\) ms for the strongest correct baseline (stock FSDP2-Muon, which runs a correct distributed Newton–Schulz on the sharded DTensor): a \(4.2\times\) optimizer-step reduction. Forward and backward stay close to the baseline: forward is \(175\) ms versus \(163\) ms (\(1.07\times\)), and backward is \(462\) ms versus \(460\) ms. This is the regime predicted by the \(\gamma_B\) model in §4.2: owner fanout is not universally better than a ring all-gather, but with balanced block owners it does not dominate the materialization path while the optimizer-step reconstruction disappears. Gather-once FSDP2-Muon is in fact slower than stock here—materializing each full matrix on every rank and re-running NS redundantly costs more than PyTorch’s sharded NS—so we headline against the faster stock baseline. This single-node \(4.2\times\) is the small end of the story: Sec. 5.5 shows the optimizer-step gap widening to \(54\times\) at 8 nodes as the baseline’s per-step matrix communication crosses the inter-node fabric while MatrixFSDP’s step stays local.
| opt-step | fwd | bwd | total | |
|---|---|---|---|---|
| stock FSDP2-Muon (torch.optim.Muon, NS on DTensor) | 367 | 163 | 460 | 991 |
| gather-once FSDP2-Muon (full matrix \(+\) local NS) | 1177 | 164 | 466 | 1808 |
| (owner-Muon) | 87 | 175 | 462 | 725 |
| speedup vs stock (strongest baseline) | 4.2\(\times\) | \(\approx\)1 | \(\approx\)1 | 1.37\(\times\) |
5.3 Convergence and numerical checks↩︎
MatrixFSDP should change data movement, not the optimizer being computed. The semantic reference is therefore full-matrix Muon, implemented as DDP Muon: every rank holds full parameters, gradients are averaged over data-parallel ranks, and Newton–Schulz runs on each complete 2D gradient. FSDP2-Muon remains the systems baseline for speed and memory, but DDP Muon is the cleaner correctness oracle because it avoids FSDP-specific reconstruction and bucket-order effects.
We check equivalence at two levels. First, a deterministic float32 harness disables TF32, fixes the decoder-LM batches, and uses distinct rank-local inputs. After each step, it compares MatrixFSDP with the full-matrix reference on losses, final logits, and reconstructed full gradients for every parameter within tolerance, directly testing that the owner rank receives the matrix gradient DDP Muon would use.
Second, both modes train the same 12-layer decoder-only LM on the same WikiText token stream for 10,000 bf16 steps (5.24B tokens), starting from identical initialization. The DDP reference uses deterministic per-parameter all-reduces rather than buckets, removing reduction-order noise. Figure 6 shows steps 3000–10,000; losses match at printed precision, not bitwise identity for arbitrary bf16 reduction orders.
5.4 Memory↩︎
Owner placement only solves the optimizer communication problem if it does not also give up ZeRO-3 memory. Table 2 compares MatrixFSDP with ZeRO-1 owner placement, where full parameters and gradients are resident on every rank. For the 4.3B model, ZeRO-1 owner placement uses \(18.5\) GB/rank; MatrixFSDP keeps the model in owner-shaped ZeRO-3 shards at \(2.7\) GB/rank, while also reducing the optimizer step from \(677\) ms to \(19\) ms. The model-size sweep in Figure 7 shows that this is not a constant offset: ZeRO-1 owner placement reaches \(54.5\) GB/rank at 12.9B and crosses the 80 GB A100 limit at \(\ge\)14B parameters, whereas MatrixFSDP uses \(6.3\) GB/rank at 14B and \(10\) GB/rank at 32B. Relative to stock FSDP2-Muon, MatrixFSDP pays only the owner’s whole-matrix residency cost (within \(\sim\)20% peak allocation here) but removes the optimizer-step gather.
| opt-step | peak alloc. mem | |
| (ms) | (GB/rank) | |
| ZeRO-1 owner placement | 677 | 18.5 |
| stock FSDP2-Muon | 651 | 2.1 |
| (ZeRO-3 owner-Muon) | 19 | 2.7 |
5.5 Scaling↩︎
MatrixFSDP’s advantage is fundamentally a scaling property (Table 3). This weak scaling grows depth with shard span; fixed-model strong scaling can reduce owner-role density and raise \(\gamma_B\) (§4.2). The baseline’s optimizer step grows from \(367\) ms on 1 node to \(5064\) ms on 8 nodes as its per-step matrix communication crosses the inter-node fabric; MatrixFSDP’s step stays flat at \(\sim\)90 ms because it issues no optimizer-step collectives. The optimizer-step speedup therefore grows \(4.2\times \to 54.6\times\) over \(1\to8\) nodes, and end-to-end step time \(1.37\times \to 2.15\times\) (peaking at 4 nodes). The end-to-end ratio is not monotone because the non-optimizer phases also scale with depth, overlap, and HSDP reduction; once MatrixFSDP’s optimizer step is flat, these phases set the residual denominator. Forward, backward, and per-rank peak allocation stay at parity with FSDP2 across the whole sweep: the HSDP backward gradient reduce (reduce-scatter over the shard group plus an all-reduce over the replicate group) is chained on the communication stream so it overlaps backward compute, exactly as FSDP2 does.
| nodes (span) | model | stock opt | opt | opt speedup | stock total | total | e2e speedup |
|---|---|---|---|---|---|---|---|
| 1 (8) | 16L / 3.2B | 367 | 87 | 4.2\(\times\) | 991 | 725 | 1.37\(\times\) |
| 2 (16) | 32L / 6.4B | 1285 | 89 | 14.5\(\times\) | 2612 | 1489 | 1.75\(\times\) |
| 4 (32) | 64L / 12.9B | 3079 | 91 | 34.0\(\times\) | 5788 | 2689 | 2.15\(\times\) |
| 8 (64) | 128L / 25.8B | 5064 | 93 | 54.6\(\times\) | 10004 | 4989 | 2.01\(\times\) |
4pt
5.6 Model-size scaling↩︎
The communication-free optimizer also holds across model size. Unlike Table 3, Table 4 fixes shard span at 64 and varies model size, measuring larger matrices at the same sharding degree. MatrixFSDP’s optimizer step is \(6\)–\(160\) ms versus \(1.0\)–\(6.2\) s for stock FSDP2-Muon (which reconstructs each matrix across all 64 ranks every step): a \(31\)–\(159\times\) optimizer-step and \(1.85\)–\(3.3\times\) end-to-end speedup, with forward and backward at parity. These ratios are not expected to be monotone: model shape and the forward/backward share change across rows, while the removed optimizer communication is only one part of total step time. Per-rank peak allocation stays at ZeRO-3 levels—MatrixFSDP is within \(\sim\)20% of plain FSDP2 and \(3\)–\(7\times\) below gather-once FSDP2-Muon, which materializes full matrices.
| model | stock opt (ms) | opt (ms) | opt speedup | e2e speedup | mem (GB) | gather-once mem (GB) |
|---|---|---|---|---|---|---|
| 1.7B | 1016 | 6.4 | 159\(\times\) | 3.3\(\times\) | 1.4 | 4.4 |
| 4B | 1235 | 14.8 | 83\(\times\) | 2.4\(\times\) | 2.6 | 9.7 |
| 8B | 1583 | 34 | 47\(\times\) | 2.0\(\times\) | 4.3 | 18.4 |
| 14B | 2286 | 73 | 31\(\times\) | 1.85\(\times\) | 6.3 | 30.7 |
| 32B | 6219 | 160 | 39\(\times\) | 2.2\(\times\) | 10.0 | 61.6 |
4pt
5.6.0.1 Load balance.
The global owner planner keeps the assignment even: across the sweep the per-rank optimizer-compute imbalance (max/avg step time) is \(1.26\)–\(1.39\) and resident-parameter-memory imbalance is \(1.18\)–\(1.21\). This is competitive with Canzona’s \(\alpha\)-balanced static partition, which reports a \(1.43\times\) load-balance ratio (improved from \(3.24\times\)), even though MatrixFSDP must also keep execution block-local for FSDP2 prefetch and overlap.
5.7 Generality beyond Muon↩︎
The owner-placement runtime is optimizer-agnostic: each 2D matrix is whole on its owner, so any single-device matrix optimizer runs there with no gather and no distributed variant. We verify this by swapping owner-Muon for Shampoo (full \(L\)/\(R\) preconditioners with inverse fourth root) and SOAP (Adam in the Shampoo eigenbasis) while keeping the same placement, collectives, routing, and checkpointing. Against a gathered reference for each optimizer, 64-GPU float32 owner-placed loss trajectories match to \(\le\) \(4\mathrm{e}{-5}\) relative for Shampoo and \(\le\) \(1.3\mathrm{e}{-4}\) for SOAP. The zero-optimizer-step-communication property is therefore inherited structurally: heavy preconditioner compute runs once per matrix on its owner instead of redundantly across ranks.
5.8 Ablations↩︎
Table 5 separates two mechanisms in the owner-Muon path: which rank owns each matrix, and how non-owner ranks materialize owner-shaped matrices during forward/backward. All rows use the same 64-layer,
12.9B-parameter model; the intended comparisons are within a row: for a fixed span, change the owner policy or collective path. Among the three custom-collective owner policies, total step time differs by at most \(3\%\). The balance column reports the per-rank optimizer-compute max/avg imbalance for the role-, scope-, and cost-aware planners (r/s/c); values closer to \(1\) mean a flatter
owner assignment. scope_greedy is flattest, while role_greedy is the memory-conservative default.
The collective implementation is not second-order. In each row, replacing the owner-segment custom collective with a matrix all-gather keeps the optimizer step local, but makes the forward/backward materialization path much more expensive. At the largest point, end-to-end time rises from \(2.51\) s to \(18.08\) s (\(7.2\times\)), with forward \(9.4\times\) and backward \(6.6\times\) slower. Thus the owner placement alone is insufficient; custom P2P owner collectives are the mechanism that makes the layout usable inside FSDP2’s normal materialize/reshard loop. In the terms of §4.2, the custom path moves the balanced owner segments, while the all-gather variant retreats to moving full matrices through the materialization path. We do not report the raw CUDA peak-allocation field from this diagnostic run as a memory result: it includes transient collective workspaces and allocator state, whereas Sec. 5.4 measures the resident ZeRO-3 memory question.
| nodes (span) | role (ms) | scope (ms) | cost (ms) | all-gather (ms) | slowdown | balance (r/s/c) |
|---|---|---|---|---|---|---|
| 1 (8) | 2833 | 2842 | 2830 | 3114 | 1.10\(\times\) | 1.07/1.02/1.07 |
| 2 (16) | 2736 | 2809 | 2729 | 5817 | 2.13\(\times\) | 1.05/1.02/1.05 |
| 4 (32) | 2565 | 2572 | 2560 | 9786 | 3.82\(\times\) | 1.12/1.03/1.11 |
| 8 (64) | 2511 | 2486 | 2508 | 18079 | 7.20\(\times\) | 1.27/1.04/1.27 |
4pt
6 Related Work↩︎
6.0.0.1 Matrix-structured optimizers.
Coordinate-wise adaptive optimizers such as Adam/AdamW [11], [12] remain the default path for non-matrix parameters, and memory-saving variants such as Adafactor [13] factor optimizer state for large matrices. Muon [1] orthogonalizes the momentum via Newton–Schulz and has shown strong token-efficiency, validated at frontier scale [2], [3]. Shampoo [4], [5] and SOAP [6] are related whole-matrix preconditioned optimizers; NorMuon [14] refines Muon with neuron-wise normalization. Unlike coordinate-wise or factored-state methods, these optimizers require a whole-matrix per-weight update, the tension this paper targets.
6.0.0.2 Distributed Muon under data parallelism.
The dominant approach reconstructs each matrix per step. Moonshot’s distributed Muon [2] uses ZeRO-1 and, per owned partition, all-gathers the full gradient, runs NS, and discards non-owned shards. In the FSDP2/ZeRO-3 world the community catalog [15], [16] enumerates all-gather-then-NS, per-rank-NS-plus-all-reduce, and gather-to-owner-then-broadcast variants. They keep equal sharding and reconstruct the matrix every step; our owner placement removes that per-step reconstruction.
6.0.0.3 Owner placement and logical/physical decoupling.
PyTorch Distributed Shampoo [9] assigns whole parameter blocks to DDP ranks via DTensor, computes the preconditioner on the owner, and all-gathers
the search direction—owner placement, but under replicated parameters. Concurrently, Canzona [10] decouples logical optimizer ownership from
physical placement for Muon/Shampoo/SOAP: its data-parallel path uses ZeRO-1 owners, and its tensor-parallel path reconstructs fragments asynchronously within the intra-node domain. MatrixFSDP is orthogonal: it realizes owner placement inside ZeRO-3
parameter sharding. This keeps ZeRO-3 per-rank memory but requires runtime pieces ZeRO-1 owner systems avoid: MatrixShard placement, block-local owner collectives, owner-buffer pinning, and owner-shard checkpointing. Tensor-parallel
coverage, including Canzona-style TP fragment reconstruction, is future/orthogonal.
6.0.0.4 Algorithmic alternatives.
Dion [17] changes the optimizer instead: a low-rank/amortized power iteration that keeps matrices sharded and synchronizes only low-rank factors—communication-friendly but not numerically equal to exact Newton–Schulz. Our approach keeps the optimizer exact and changes the layout.
6.0.0.5 Distributed training systems.
Synchronous data-parallel runtimes such as Horovod and PyTorch DDP [18], [19] make replicated-parameter training fast through efficient all-reduce, bucketing, and communication/computation overlap. Tensor, pipeline, and compiler-planned parallel systems [20]–[27] split computation or tensors across devices and are complementary to data-parallel sharding. ZeRO and FSDP/FSDP2 [7], [8], [16], [28], [29] provide the sharded-parameter substrate; we extend their equal-sharding placement with whole-parameter owner roles for matrices and non-matrix tails, expressed as a custom DTensor placement.
7 Conclusion↩︎
MatrixFSDP is useful when sharding makes the optimizer step, rather than the model layers, the expensive part of training. It keeps each matrix whole only at its assigned owner, so Muon-style updates run locally while forward, backward, and checkpointing still follow the ZeRO-3 discipline.
This is the strongest fit for multi-node ZeRO-3 runs with large shard spans, where reconstructing matrices at every optimizer step would cross the inter-node fabric. The gain is smaller on single-node runs, when high gradient accumulation amortizes one optimizer step over many forward/backward passes, or under fixed-model strong scaling where \(W\) grows far beyond the number of sizable owner roles in a block. In that last regime, \(\gamma_B\) can grow and owner fanout can become the next bottleneck, making hierarchical owner-fanout a natural extension.