The reviewed record of science sign in
Pith

arxiv: 2607.05895 · v1 · pith:DJPV7COV · submitted 2026-07-07 · cs.DC

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

Ming Gao , Yanwu Xu , Hao Zhang This is my paper

Reviewed by Pith T0 review T1 audit T2 compute T3 formal T4 kernel 2026-07-08 21:10 UTCglm-5.2pith:DJPV7COVrecord.jsonopen to challenge →

classification cs.DC
keywords ZeRO-3FSDPMuonmatrix optimizerNewton-Schulzparameter shardingowner placementdistributed training
0
0 comments X

The pith

Zero optimizer-step communication under ZeRO-3

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

MatrixFSDP eliminates the per-step matrix communication that whole-matrix optimizers like Muon require under ZeRO-3 parameter sharding. The key insight is a placement change: instead of slicing every weight equally across all data-parallel ranks, assign each 2D weight matrix to a single owner rank that holds the complete matrix while other ranks hold empty shards. The standard backward gradient reduction then deposits the full matrix gradient on the owner, so Newton-Schulz orthogonalization runs locally with no gather, broadcast, or redistribute. Non-matrix tensors (LayerNorm, biases) are packed into tail owners that stay on AdamW. The system preserves ZeRO-3 per-rank memory because only one copy of each parameter exists across the data-parallel group; it simply lives on one rank rather than being equally sliced. The runtime challenge is that owner-shaped shards break FSDP2's equal-size collective assumptions, requiring a global balance-aware owner planner, deterministic peer-to-peer owner-segment collectives that move only nonempty pieces, owner-buffer pinning to handle storage lifecycle hazards, and owner-shard checkpoint resharding. The resulting optimizer trajectory matches full-matrix Muon (fp32 tolerance, bf16 printed loss match over 10K steps) while achieving up to 54.6x optimizer-step speedup at 8 nodes and up to 2.15x end-to-end speedup, with per-rank memory under 10 GB at 32B parameters.

Core claim

The central mechanism is owner-shaped ZeRO-3 placement: changing where shards live (one rank owns the whole matrix, others hold empty shards) rather than changing the optimizer or the sharding strategy. This placement exploits the fact that ZeRO-3's backward gradient reduction already aggregates gradients into the local shard shape, so if the local shard is the full matrix on the owner, the optimizer input arrives for free. The owner planner balances resident bytes and optimizer work across ranks using a greedy assignment over matrix and tail roles, while custom P2P collectives move only nonempty owner segments during forward/backward materialization, avoiding both padding waste and full-mat

What carries the argument

MatrixShard metadata, balance-aware owner planner, owner-segment P2P collectives, owner-buffer pinning, owner-shard checkpoint resharding

If this is right

  • Whole-matrix optimizers (Muon, Shampoo, SOAP) become practical for multi-node ZeRO-3 training at scales where ZeRO-1 owner placement exceeds GPU memory, removing a key systems obstacle to adopting these token-efficient optimizers at frontier scale.
  • The owner-placement approach is structurally optimizer-agnostic: any single-device matrix optimizer runs on the owner with zero communication, so future matrix-structured optimizers inherit the communication-free property without new distributed implementations.
  • The scaling advantage grows with node count because the baseline's per-step matrix communication crosses inter-node fabric while MatrixFSDP's step stays local, making the method increasingly valuable as training clusters grow.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The bandwidth model shows owner fanout cost is proportional to the imbalance factor gamma_B; when one rank owns most of a block (gamma_B approaches W), the owner fanout can exceed ordinary FSDP's ring all-gather. This suggests a degradation regime under fixed-model strong scaling where the shard span W grows far beyond the number of sizable owner roles, which the paper acknowledges but does not fu
  • Hierarchical owner-fanout, where owner roles are nested across node boundaries, could extend the approach to very large shard spans by reducing per-owner fanout degree, but this is left as future work.
  • The interaction between owner placement and tensor parallelism is unexplored: matrices fragmented by TP are skipped by the owner planner, so models using hybrid TP+DP sharding would need a complementary mechanism for the TP-fragmented weights.

Load-bearing premise

The global owner planner can always find a balanced assignment where no single rank owns disproportionately more bytes than others, so that the peer-to-peer owner fanout during forward/backward materialization does not become a new communication bottleneck. When the shard span grows far beyond the number of available owner roles in a block, this balance can degrade and owner fanout cost can approach or exceed ordinary FSDP's collective cost.

What would settle it

If the owner planner cannot maintain balance (gamma_B near 1) as shard span grows, forward/backward materialization through owner-segment P2P collectives should become slower than ordinary FSDP2's ring all-gather, eliminating the end-to-end speedup even as the optimizer step itself remains communication-free.

Figures

Figures reproduced from arXiv: 2607.05895 by Hao Zhang, Ming Gao, Yanwu Xu.

Figure 1
Figure 1. Figure 1: Three ways to combine sharding with a whole-matrix optimizer. (a) Keeping ordinary ZeRO-3 sharding preserves memory but reconstructs each 2D gradient at every optimizer step. (b) ZeRO-1 owner placement makes the optimizer local, but full parameters are resident on every rank. (c) MatrixFSDP makes owner placement a ZeRO-3 layout: backward already leaves the full gradient on the owner, while forward/backward… view at source ↗
Figure 2
Figure 2. Figure 2: Owner planning and runtime materialization use different granularities. The planner balances owners over all wrapped blocks and emits one MatrixShard layout per block; the runtime still gathers, prefetches, computes, and reshards one block at a time. Algorithm 1 Balance-aware owner planning 1: Build roles: each 2D matrix is one role; non-2D tensors form one tail role per block. 2: ROLEGREEDY: per block, as… view at source ↗
Figure 3
Figure 3. Figure 3: Owner-buffer pinning within FSDP2’s unshard/reshard lifecycle. Forward materializes the full matrix and autograd may save views into that storage. Reshard shrinks the buffer but pins the same allocation instead of returning it to the pool; backward revives the full matrix in that same storage, and the owner releases it only after Newton–Schulz. 4.4 Optimizer and checkpoint state After backward, every FSDP … view at source ↗
Figure 4
Figure 4. Figure 4: Local shards after reshard. Ordinary FSDP2 slices every tensor across ranks; MatrixFSDP keeps each 2D matrix and the non-2D tail as whole owner roles, with empty shards on the other ranks. 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); c… view at source ↗
Figure 5
Figure 5. Figure 5: Real-data convergence on WikiText after step 3000. Ma￾trixFSDP overlays a DDP Muon reference (full parameters, aver￾aged full gradients, whole-matrix Muon); the full 10,000-step log records max printed |∆loss| = 0. ness oracle because it avoids FSDP-specific reconstruction and bucket-order effects. We check equivalence at two levels. First, a determinis￾tic float32 harness disables TF32, fixes the decoder-… view at source ↗
Figure 7
Figure 7. Figure 7: Optimizer-step latency vs node count (weak scaling, depth scaled with shard span, 8–64 GPUs, InfiniBand; log scale). The strongest correct baseline (stock FSDP2-Muon) grows 367 → 5064 ms as its per-step Newton–Schulz matrix communication crosses the inter-node fabric, while MatrixFSDP’s communication￾free step stays flat (∼90 ms). The optimizer-step speedup (la￾bels) therefore grows with scale, 4.2× → 54.6… view at source ↗
read the original 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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

1 major / 6 minor

Summary. MatrixFSDP proposes assigning each 2D weight matrix to a single owner rank within ZeRO-3 parameter sharding, so that whole-matrix optimizers (Muon, Shampoo, SOAP) can run locally without optimizer-step communication. The backward gradient reduction deposits the full matrix gradient on the owner rank, Newton-Schulz runs there, and no post-backward matrix all-gather is needed. The system contributions include a global owner planner, MatrixShard metadata, owner-segment P2P collectives, owner-buffer pinning, and checkpoint resharding. The evaluation verifies numerical correctness against DDP Muon (fp32 tolerance, bf16 loss match over 10K steps), measures ZeRO-3-scale memory (under 10 GB/rank at 32B parameters), and reports up to 54.6x optimizer-step speedup and 2.15x end-to-end speedup at 8 nodes.

Significance. The paper addresses a real and timely systems problem: the mismatch between ZeRO-3 sharding and whole-matrix optimizers. The core idea of owner placement inside ZeRO-3 (as opposed to ZeRO-1 owner placement with full parameters resident) is novel and well-motivated. Strengths include: (1) a first-principles bandwidth model (§4.2) that predicts when owner fanout is competitive with ring all-gather; (2) thorough correctness verification at fp32 tolerance against a DDP Muon reference over 10K bf16 training steps (§5.3); (3) memory measurements showing ZeRO-3-scale per-rank allocation up to 32B parameters (§5.4); (4) weak-scaling and fixed-span sweeps covering 1-8 nodes and 1.7B-32B models (§5.5, §5.6); (5) generality beyond Muon demonstrated for Shampoo and SOAP (§5.7); (6) ablations isolating the contribution of custom collectives vs. all-gather (§5.8). The approach is falsifiable: the bandwidth model gives a concrete condition (gamma_B near 1) under which the method is efficient, and the evaluation tests the complementary regime.

major comments (1)
  1. §5.5, Table 3 and §4.2: The weak-scaling sweep scales model depth with shard span, keeping the ratio of owner roles to ranks roughly constant. The paper's own bandwidth model shows that when gamma_B grows (approaching W), owner fanout cost reaches (W-1)/W * P_B, which can exceed ordinary FSDP's ring all-gather. The conclusion (§7) acknowledges that fixed-model strong scaling where W grows far beyond the number of owner roles causes gamma_B to grow, but no experiment measures this regime. Fixed-model strong scaling—where a user has a fixed model and increases GPU count—is a common practical deployment scenario. The paper should either (a) add an experiment measuring forward/backward latency degradation as W grows beyond owner-role count for a fixed model, or (b) more precisely bound the regime where the approach breaks even with stock FSDP2-Muon, giving the reader a concrete crossover W.
minor comments (6)
  1. Title: 'communication-free' is slightly imprecise since forward/backward still communicate via owner-segment P2P. Consider 'optimizer-step communication-free' for precision.
  2. §4.2: The bandwidth model defines gamma_B, C_FSDP, and C_owner but does not explicitly state the units of C_FSDP and C_owner (bytes per step). This is inferrable from context but would benefit from stating it directly.
  3. Table 5: The balance column header '(r/s/c)' is defined in the caption but could also be stated in the column header for clarity (r=role_greedy, s=scope_greedy, c=costaware).
  4. §5.3, Figure 5: The convergence plot shows steps 3000-100,000. The caption mentions 'max printed |delta_loss| = 0' but the y-axis range (5-5.8) makes it hard to visually confirm overlap. Consider showing earlier steps or a difference plot.
  5. §5.8: The paper notes that raw CUDA peak allocation from the ablation run is not reported because it includes transient workspaces. It would help to state the magnitude of this overhead to reassure readers it does not undermine the memory claims in §5.4.
  6. §6: The comparison with Canzona (Wang et al., 2026) is discussed in terms of ZeRO-1 vs. ZeRO-3, but a brief note on whether Canzona's alpha-balanced partitioning could improve MatrixFSDP's owner planner would preempt reader questions.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for the careful reading and the constructive recommendation. The referee correctly identifies a gap between our bandwidth model (§4.2) and our experiments: we acknowledge in §7 that fixed-model strong scaling can cause γ_B to grow, but we do not experimentally measure that regime or give a concrete crossover W. We agree this should be addressed and will add both an experiment and a sharper bound in the revision.

read point-by-point responses

  1. Referee: §5.5, Table 3 and §4.2: The weak-scaling sweep scales model depth with shard span, keeping the ratio of owner roles to ranks roughly constant. The paper's own bandwidth model shows that when gamma_B grows (approaching W), owner fanout cost reaches (W-1)/W * P_B, which can exceed ordinary FSDP's ring all-gather. The conclusion (§7) acknowledges that fixed-model strong scaling where W grows far beyond the number of owner roles causes gamma_B to grow, but no experiment measures this regime. Fixed-model strong scaling—where a user has a fixed model and increases GPU count—is a common practical deployment scenario. The paper should either (a) add an experiment measuring forward/backward latency degradation as W grows beyond owner-role count for a fixed model, or (b) more precisely bound the regime where the approach breaks even with stock FSDP2-Muon, giving the reader a concrete crossover W.

    Authors: The referee is correct that our evaluation covers the weak-scaling regime (depth grows with shard span, keeping owner-role density roughly constant) but does not experimentally test the fixed-model strong-scaling regime where W grows beyond the number of sizable owner roles in a block. We agree this is an important practical scenario and that the paper should address it more concretely than the brief acknowledgment in §7. We will make two changes in the revision. First, we will add a fixed-model strong-scaling experiment: holding a 12.9B-parameter model constant and increasing shard span from 8 to 64, we will report forward, backward, and optimizer-step latency for both MatrixFSDP and stock FSDP2-Muon, showing how γ_B grows and where the forward/backward materialization cost begins to offset the optimizer-step savings. Second, we will sharpen the break-even analysis from §4.2 into a concrete crossover condition. The bandwidth model gives C_owner_B = γ_B · (W−1)/W · P_B versus C_FSDP_B ≈ (W−1)/W · P_B for the materialization path, while the optimizer-step reconstruction that MatrixFSDP removes costs roughly (W−1)/W · P_B per step on the critical path. MatrixFSDP is net beneficial when the incremental forward/backward fanout penalty (γ_B − 1) · (W−1)/W · P_B is smaller than the removed optimizer-step reconstruction. For a transformer block with M 2D weight matrices and shard span W, the number of owner roles per block is at most 2M+1 (including the tail); when W exceeds this, γ_B begins to grow. We will state this crossover explicitly and verify it against the new experiment. We believe this addresses both options (a) and (b) the referee proposed. revision: yes

Circularity Check

0 steps flagged

No circularity found: systems paper with first-principles bandwidth model and independent empirical validation

full rationale

This is a systems paper, not a theoretical derivation with fitted constants. The bandwidth model in §4.2 (γ_B, C_FSDP, C_owner) is a straightforward first-principles derivation from communication volumes — each formula follows directly from counting bytes moved per rank. No parameter is fitted to data and then presented as a prediction. The core design argument (owner placement → backward reduction deposits full gradient on owner → Newton-Schulz runs locally → no optimizer-step collective) is a logical consequence of the system layout, not a tautology: the owner holds the full shard by construction, and the standard FSDP gradient reduction naturally lands the full gradient there. The correctness claim is validated against an independent DDP Muon reference (fp32 tolerance, bf16 loss match over 10K steps). The speedup and memory claims are validated by independent latency and memory measurements (Tables 1–4, Figure 6). The authors do not cite their own prior work as load-bearing evidence. The skeptic's concern about the evaluation avoiding the fixed-model strong-scaling regime where γ_B grows is a valid evaluation-completeness concern, but it is not circularity: the paper explicitly acknowledges this regime in both §5.5 and the conclusion, and the bandwidth model that predicts the degradation is derived independently of the measurements. No step in the derivation chain reduces to its own inputs by construction.

Axiom & Free-Parameter Ledger

0 free parameters · 3 axioms · 2 invented entities

No free parameters are fitted to data. The planner algorithms (ROLE_GREEDY, SCOPE_GREEDY, COSTAWARE) are design choices evaluated empirically, not fitted constants. No invented physical entities — all entities are software components with measured performance.

axioms (3)
  • domain assumption Backward gradient reduction under ZeRO-3 deposits the full 2D gradient on the owner rank when that rank holds the whole matrix shard.
    §3: '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.' This follows from ZeRO-3 reduction semantics but is the structural foundation of the approach.
  • domain assumption FSDP2's materialize/reshard lifecycle can be preserved with owner-shaped (unequal-size) shards by replacing equal-size collectives with owner-segment P2P collectives.
    §4: The entire runtime design (§4.1–4.4) assumes that the unshard/reshard/prefetch schedule remains valid when shards are unequal. The ablation in Table 5 confirms this is non-trivial — all-gather fallback is 7.2x slower — but the assumption that the custom path works is load-bearing.
  • domain assumption The global owner planner can find assignments where γB ≈ 1 for the model architectures and shard spans tested.
    §4.2 and §5.8: The bandwidth model shows owner fanout cost is proportional to γB. The evaluation reports imbalance ratios of 1.18–1.39 (§5.6), but the paper does not prove this holds for arbitrary model architectures or shard spans. The conclusion acknowledges degradation when W grows beyond the number of sizable owner roles.
invented entities (2)
  • MatrixShard metadata independent evidence
    purpose: Carries owner placement information (which rank owns each matrix/tail, segment sizes) through optimizer routing, buffer pinning, checkpointing, and cross-world-size resharding.
    MatrixShard is a system data structure, not a physical entity. Its correctness is verified by the fp32 equivalence tests (§5.3) and the checkpoint resharding path. It is not postulated to explain a phenomenon; it is an implementation component.
  • Owner-segment P2P collectives independent evidence
    purpose: Replace equal-size FSDP2 collectives with deterministic send/recv schedules that move only non-empty owner segments during materialization and gradient reduction.
    System component validated by latency measurements (Tables 1, 3, 4) and ablation (Table 5). Not a postulated entity.

pith-pipeline@v1.1.0-glm · 16229 in / 4339 out tokens · 252049 ms · 2026-07-08T21:10:16.823384+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

14 extracted references · 14 canonical work pages · 10 internal anchors

  1. [1]

    Ahn, K. et al. Dion: Distributed orthonormalized updates. arXiv preprint arXiv:2504.05295,

    work page arXiv

  2. [2]

    Scalable Second Order Optimization for Deep Learning

    Anil, R., Gupta, V ., Koren, T., Regan, K., and Singer, Y . Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018,

    work page internal anchor Pith review Pith/arXiv arXiv 2002

  3. [3]

    Practical Efficiency of Muon for Pretraining

    Essential AI. Practical efficiency of muon for pretraining. arXiv preprint arXiv:2505.02222,

    work page internal anchor Pith review Pith/arXiv arXiv

  4. [4]

    NorMuon: Making muon more efficient and scalable.arXiv preprint arXiv:2510.05491,

    Li, Z., Liu, L., Liang, C., Chen, W., and Zhao, T. NorMuon: Making muon more efficient and scalable.arXiv preprint arXiv:2510.05491,

    work page arXiv

  5. [5]

    TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training

    Liang, W., Liu, T., et al. TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training. arXiv preprint arXiv:2410.06511,

    work page internal anchor Pith review Pith/arXiv arXiv

  6. [6]

    Liu, J. et al. Muon is scalable for LLM training.arXiv preprint arXiv:2502.16982,

    work page internal anchor Pith review Pith/arXiv arXiv

  7. [7]

    ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

    Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y . ZeRO: Memory optimizations toward training trillion parameter models.arXiv preprint arXiv:1910.02054,

    work page internal anchor Pith review Pith/arXiv arXiv 1910

  8. [8]

    A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale

    Shi, H.-J. M., Lee, T.-H., Iwasaki, S., Gadekar, J., et al. A distributed data-parallel PyTorch implementation of the distributed shampoo optimizer for training neural net- works at-scale.arXiv preprint arXiv:2309.06497,

    work page internal anchor Pith review Pith/arXiv arXiv

  9. [9]

    Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

    Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LM: Training multi- billion parameter language models using model paral- lelism.arXiv preprint arXiv:1909.08053,

    work page internal anchor Pith review Pith/arXiv arXiv 1909

  10. [10]

    SOAP: Improving and Stabilizing Shampoo using Adam

    Vyas, N., Morwani, D., et al. SOAP: Improving and stabilizing shampoo using adam.arXiv preprint arXiv:2409.11321,

    work page internal anchor Pith review Pith/arXiv arXiv

  11. [11]

    Canzona: A uni- fied, asynchronous, and load-balanced framework for distributed matrix-based optimizers.arXiv preprint arXiv:2602.06079,

    Wang, L., Men, R., Liu, D., et al. Canzona: A uni- fied, asynchronous, and load-balanced framework for distributed matrix-based optimizers.arXiv preprint arXiv:2602.06079,

    work page arXiv

  12. [12]

    GSPMD: General and Scalable Parallelization for ML Computation Graphs

    Xu, Y ., Lee, H., Chen, D., Hechtman, B., Huang, Y ., Joshi, R., Krikun, M., Lepikhin, D., Ly, A., Maggioni, M., Pang, R., Shazeer, N., Wang, S., Wang, T., Wu, Y ., and Chen, Z. GSPMD: General and scalable parallelization for ML computation graphs.arXiv preprint arXiv:2105.04663,

    work page internal anchor Pith review Pith/arXiv arXiv

  13. [13]

    PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

    Zhao, Y ., Gu, A., Varma, R., et al. PyTorch FSDP: Ex- periences on scaling fully sharded data parallel.arXiv preprint arXiv:2304.11277,

    work page internal anchor Pith review Pith/arXiv arXiv

  14. [14]

    H., Haj-Ali, A., Wang, Y ., Yang, J., Zhuo, D., Sen, K., Gonzalez, J

    Zheng, L., Jia, C., Sun, M., Wu, Z., Yu, C. H., Haj-Ali, A., Wang, Y ., Yang, J., Zhuo, D., Sen, K., Gonzalez, J. E., and Stoica, I. Alpa: Automating inter- and intra- operator parallelism for distributed deep learning. In USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2022

    work page 2022