Training language models whose parameters and optimizer state exceed a single node's memory requires sharding these tensors across devices. Under tensor parallelism (TP) or fully sharded data parallelism (FSDP), matrix-aware optimizers such as Muon introduce communication to perform global matrix operations across shards. We propose Dobb (Distributed Optimizer with Block-diagonal Eigenbases), a matrix-aware optimizer designed to compute updates locally within parameter shards. Dobb approximates the input-activation second moment with diagonal blocks and tracks a complete basis within each block. It projects momentum into these bases, applies sign normalization to the coefficients, and reconstructs the update. When each block lies within one shard under TP or FSDP, Dobb applies its matrix update locally, requiring no communication beyond data-parallel synchronization of the activation second moments. We benchmark Dobb against Dion, a distributed matrix-aware optimizer designed to retain the benefits of orthonormalized updates while reducing communication. On GPT-2 XL (1.56B parameters), Dobb with 16 blocks reaches a validation loss of 2.476, compared with 2.511 for Dion at rank fraction 1/16. Under FSDP on 32 H100 GPUs, Dobb reduces the amortized optimizer-call time from 124.1 ms for Dion to 9.2 ms.
Figures & tables
Figure 1: GPT-2 XL (1.56B parameters) validation loss and optimizer-specific cross-shard payload. Losses are measured under the shared update-scale protocol (Table 4 ). Labels give block count B or rank fraction ρ . For a weight matrix of size m×n , optimizer-specific cross-shard payload is the logical number of scalar entries exchanged between devices holding different shards of that matrix during one optimizer step; collective factors and bytes per entry are omitted. We express this derived count relative to Muon’s dense count mn . At Dion rank r , the per-matrix ratio is r(m+n)/(mn) ; the aggregate ratio for the plotted GPT-2 XL matrices is 1.5ρ . Payloads are derived for FSDP with FS=4 , where each Dobb block lies within one shard; Dobb and AdamW then have zero payload under this definition. The count excludes forward-backward communication and Dobb’s separately counted data-parallel synchronization of activation second moments (Section 5 ).
Figure 2: Effective rank of the second moments from GPT-2 Small runs on OpenWebText. Within each plot, four panels show transformer layer depth at the indicated fraction of training completed; curves distinguish weight-matrix types. From 50% to 100% of training, the effective-rank distribution remains stable. Effective rank is exp(−∑ipilogpi) as a percentage of the factor dimension, n for At and m for Gt , where pi is the i th normalized eigenvalue. The second-moment matrices are averaged over accumulated mini-batches. The output-gradient plot separates the query, key, and value blocks of the fused attention projection.
(a) Sign of a matrix
Orthogonality
Effective rank of Y
Y
sign(Y)
80%
0.736
0.775
50%
0.412
0.643
20%
0.205
0.543
5%
0.109
0.457
Table 1: Orthogonality of updates on synthetic 768×768 gradients. (a) A matrix Y with prescribed effective rank and its sign. (b) Sign normalization (Dobb) and column normalization in the same eigenbasis of the full activation second moment (a single block), on gradients DX⊤ whose output-gradient second moment G has the stated effective rank, with the effective rank of the activation second moment fixed at 15%. Effective rank is the entropy-based measure of Figure 2 as a percentage of the dimension, computed from the singular values of Y in (a). Orthogonality is ∥O∥∗/(min(m,n)∥O∥F) for a nonzero m×n matrix O , where ∥⋅∥∗ and ∥⋅∥F denote the nuclear and Frobenius norms; it equals 1 exactly when all singular values are equal. Appendix D describes the construction.
Dobb blocks B
Dobb orthogonality
Dion rank r
Dion orthogonality
1
0.849
768
0.804
4
0.834
192
0.427
16
0.815
48
0.208
Table 2: Update orthogonality under block partitioning and rank truncation. Orthogonality is the metric defined in Table 1 . Values average 72 weight matrices using the same momentum from a GPT-2 Small checkpoint at optimizer step 5,000 with second-moment capture at every step.
Method
Optimizer-specific cross-shard payload
Muon, dense exchange
O(mn)
Dion, rank r
O(r(m+n))
Dobb
0
Table 3: Leading optimizer-specific cross-shard payload for one m×n matrix. The counts apply when either the input or output axis is sharded, under TP or FSDP. For Dobb, each block lies within one shard. Constant collective factors and element sizes are omitted. The counts exclude the data-parallel synchronization of Dobb’s block second moments, n2/B entries per matrix at each refresh.
Optimizer
Blocks / rank fraction
Small
Medium
Large
XL
AdamW
n/a
3.044
2.761
2.606
2.497
Muon
full
3.004
2.725
2.574
2.460
RMNP
n/a
3.083
2.778
2.618
2.499
Dion
1
3.002
2.720
2.568
2.453
1/4
3.038
2.754
2.596
2.481
1/16
3.108
2.802
2.636
2.511 †
Table 4: Final validation loss under the shared update-scale protocol, with common learning rate and weight decay within each model scale. Small, Medium, Large, and XL train for 10,000, 20,000, 30,000, and 40,000 steps. Dion ranks are fractions of each weight matrix’s full rank.
Optimizer call (ms)
Iteration (ms)
Model
GPU
DP
FS
TP
GPUs
AdamW
Muon
Dion
Dobb
AdamW
Muon
Dion
Dobb
GPT-2 Large
A5000
1
10
2
20
11.3
623.0
1224.1
6.9
7764
8376
8977
7760
GPT-2 XL
A5000
2
8
2
32
18.7
946.5
2523.2
16.3
11,649
12,577
14,153
11,646
GPT-2 XL
H100
2
16
1
32
9.3
82.7
232.7
9.2
1069
1143
1293
1069
GPT-2 3B
H100
1
16
2
32
49.4
454.3
1385.6
7.6
2388
2793
3724
2346
GPT-2 6B
H100
1
8
4
32
13.9
2106.8
1505.7
10.9
4823
6916
6315
4820
Table 5: Times in milliseconds, with DP data-parallel replicas and FSDP shard degree FS . GPT-2 3B, 6B, and 8B are scaled-up GPT-2-style models. Optimizer call is the mean duration of optimizer.step() ; Dobb’s is averaged over one refresh cycle, (Refresh+8⋅Non-refresh)/9 . Iteration is the pooled iteration time defined in Appendix G.1 . Dion uses rank r=400 on Large; rank fraction ρ=1/3 on XL (A5000), 6B, and 8B; ρ=1/4 on XL (H100) and 3B; and on Qwen3, a rank fraction matched to the block width 1024/FS (Appendix G.5 ). For GPT-2 XL on H100, Table 23 also reports Dion at ρ=1/16 , the rank fraction of the loss comparison in the abstract, with an optimizer call of 124.1 ms. Dashes mark methods not run on Qwen3. Appendix G gives Dobb’s step() times on refresh and non-refresh iterations and all other configurations.
Appendix figures & tables21 assets
Supplementary material from the paper’s appendix.
Appendix
Symbol
Meaning
Defined in
Layer and gradient statistics
m , n
Output and input dimensions of a weight matrix
Section 3
Wt , t
Weight matrix of a layer y=Wx at optimizer step t
Section 3
Bt , N
Mini-batch at step t , and its number of examples or token positions
Section 3
fBt
Mean loss on Bt
Equation 1
xi , gi
Input activation and output gradient of the i th contribution
Equation 1
Appendix
Table 6: Notation. Parenthesized values are the principal settings.
Configuration
Output axis m sharded
Input axis n sharded
Dobb
0
0
Dobb-G
0
0
Dobb- ℓ2
O(n)
0
Dobb-G- ℓ2
0
O(m)
Appendix
Table 7: Leading optimizer-specific cross-shard payload for the basis and normalization ablations on one m×n matrix, with each block within one shard. The ℓ2 entries are scalars reduced per update; constant collective factors and element sizes are omitted. The counts exclude the data-parallel synchronization of the block second moments.
Field
Small
Medium
Large
XL
Selected learning rate
3.54×10−3
3.54×10−3
2.5×10−3
2.5×10−3
Grid index j
−1
−1
−2
−2
Evaluation interval (steps)
1,000
1,000
1,000
1,000
Evaluation sequences
3,000
3,000
3,000
3,000
Seeds
5000
5000
5000
5000
Appendix
Table 8: Training protocol for the principal quality comparison. Table 10 gives the architecture, step count, training token budget, and AdamW grid at each size. Grid indices j refer to Equation 7 . Validation loss is computed on a fixed set of held-out sequences.
Raw update
∥O∥F2
s
RMS of sO
GPT-2
Full-rank orthogonalized Muon
min(m,n)
τmax(m,n)
τ
0.2
Dion, r orthonormal left / unit-norm right columns
r
τmax(m,n)
τr/min(m,n)
0.2ρ
RMNP, unit-norm rows
m
τmax(m,n)
τmax(m,n)/n
0.2 or 0.4 a
Dobb, Dobb- ℓ2
n
τmax(m,n)
τmax(m,n)/m
0.2 or 0.4 b
Dobb-G, Dobb-G- ℓ2
m
τmax(m,n)
τmax(m,n)/n
0.2 or 0.4 a
Appendix
Table 9: Raw-update norm, update scale s , and resulting update RMS of each matrix optimizer in exact arithmetic with nonzero coefficients. Approximate orthogonalization and zero columns can lower or perturb the attained RMS. The GPT-2 column uses τ=0.2 and the rank fraction ρ=r/min(m,n) . Every entry of the GPT-2 column lies in the range of 0.2 to 0.4 that Liu et al. [2025] observe for AdamW, except Dion’s at ρ<1 .
Nparam
Layers
Width
Heads
Steps
D
Grid j
Runs
γ⋆ ( 10−3 )
Swept sizes
62M
6
640
10
5,000
2.46B
−4…4
9+5a
4.16 [3.74, 4.45]
124M
12
768
12
10,000
4.92B
−4…4
9
3.56 [3.42, 3.72]
219M
18
896
14
15,000
7.37B
−3…1
5
3.56 [3.31, 3.83]
354M
24
1024
16
20,000
9.83B
−3…1
5
3.13 [2.43, 3.60]
Extrapolated sizes
Appendix
Table 10: Model sizes and learning-rate sweeps. Nparam is the parameter count, including token embeddings and excluding position embeddings, and D is the training token budget. Sweeps use the grid of Equation 7 , and grid ranges are inclusive. For swept sizes, γ⋆ is the optimum of a quadratic in j fitted to the five grid points centered on the sweep’s best run, with both seeds at 62M; its 68% interval comes from a parametric bootstrap at the measured seed noise. For 772M and 1.56B, γ⋆ is the extrapolation of Equation 11 with a 68% residual-bootstrap interval; both models are trained at γ=2.5×10−3 ( j=−2 ).
Figure 3: Learning-rate scaling. (a) Final validation loss relative to the fitted per-size minimum, with the shared-curvature parabolas of Equation 8 ; ticks at zero mark the fitted optima. Both seeds at 62M are shown, and the bar marks the estimated run-to-run standard deviation. (b) Optimal learning rate against parameter count. Filled points are the per-size optima of Table 10 with 68% intervals, and the line and band are the fitted law of Equation 11 with its 68% bootstrap band. Each 68% interval and band spans the 16th to 84th percentiles of its bootstrap distribution, which equals ±1 standard deviation for a normal distribution. The open circles are the extrapolated optima at 772M and 1.56B, and the diamonds mark the learning rate used for both, γ=2.5×10−3 ( j=−2 ). The shaded region lies beyond the largest swept size.
Small
Medium
Large
XL
B
Dobb
Dobb- ℓ2
Δ
Dobb
Dobb- ℓ2
Δ
Dobb
Dobb- ℓ2
Δ
Dobb
Dobb- ℓ2
Δ
1
3.017
3.043
0.026
2.735
2.753
0.018
2.579
2.599
0.020
2.461
2.476
0.015
4
3.024
3.037
0.013
2.741
2.760
0.019
2.588
2.609
0.021
2.471
2.483
0.012
16
3.026
3.036
0.010
2.746
2.760
0.014
2.595
2.615
0.020
2.476
2.491
0.015
64
3.037
3.045
0.008
2.742
2.759
0.017
2.600
2.619
0.019
2.482
2.498
0.016
Appendix
Table 11: Normalization ablation for the activation-side basis: final validation loss for Dobb (sign) and Dobb- ℓ2 (column ℓ2 ). The Dobb columns repeat Table 4 ; Δ is the Dobb- ℓ2 loss minus the Dobb loss.
Activation basis
Output-gradient basis
B
Dobb
Dobb- ℓ2
Dobb-G
Dobb-G- ℓ2
1
3.017
3.043
3.071
3.077
4
3.024
3.037
3.067
3.071
16
3.026
3.036
3.076
3.075
64
3.037
3.045
3.072
3.083
Appendix
Table 12: Basis-source and normalization ablation. The activation-side columns repeat the Small entries of Table 11 .
Norm
Source
Final loss
ℓ∞
Activation second moment (Dobb)
3.017
Gradient Gram, input side
3.027
Gradient Gram, output side
3.055
Output-gradient second moment (Dobb-G)
3.071
ℓ2
Activation second moment (Dobb- ℓ2 )
3.043
Gradient Gram, input side
3.042
Appendix
Table 13: Second-moment-source ablation.
T
Final loss
Full iteration (ms)
1
3.005
762.5
4
3.014
459.3
9
3.017
403.6
16
3.017
384.8
100
3.015
363.4
Appendix
Table 14: Refresh-period ablation for Dobb. Full iteration reports the amortized time.
Full iteration (ms)
Basis refresh
Final loss
Non-refresh
Refresh
Amortized
k=1 (principal)
3.017
359.5
756.5
403.6
k=2
3.019
359.9
1137.8
446.4
k=3
3.017
358.5
1512.3
486.7
Exact eigendecomposition
3.015
360.1
1388.9
474.4
Appendix
Table 15: Orthogonal-iteration ablation for Dobb.
Full iteration (ms)
B
QR routine
Final loss
Non-refresh
Refresh
Amortized
1
Householder
3.017
365.2
760.3
409.1
Cholesky
3.016
361.7
486.3
375.6
4
Householder
3.024
394.2
604.7
417.6
Cholesky
3.022
393.2
431.8
397.5
Appendix
Table 16: QR-routine ablation for Dobb.
Full iteration (ms)
Method
Capture policy
Final loss
Non-refresh
Refresh
Amortized
Dobb
First micro-batch
3.017
360.9
755.7
404.8
All micro-batches
3.021
361.8
888.1
420.3
Dobb-G
First micro-batch
3.071
370.4
835.1
422.0
All micro-batches
3.071
370.2
894.7
428.5
Appendix
Table 17: Capture-policy ablation for Dobb and Dobb-G.
Accumulation
Basis
Final loss
Median iteration (ms)
Storage per matrix (bytes)
fp32
fp32
3.023
361.2
8n2/B
fp32
bf16 (principal)
3.017
360.4
6n2/B
bf16
fp32
3.028
361.0
6n2/B
bf16
bf16
3.027
360.9
4n2/B
Appendix
Table 18: Precision ablation for Dobb. Median iteration is the median full-iteration time in milliseconds on iterations without an eigenbasis refresh. The storage column gives the second-moment and basis storage per matrix and is derived from buffer sizes, not measured.
Method
B
Off
On
Dobb
1
3.017
3.014
Dobb
4
3.024
3.023
Dobb
16
3.026
3.030
Muon
n/a
3.004
3.004
RMNP
n/a
3.083
3.072
Appendix
Table 19: Nesterov ablation.
Optimizer
Mean
Run-to-run SD
Range width
AdamW
7723
22.8
77
Muon
7788
26.9
81
Dion r=400
7713
16.0
56
Dobb
7740
15.3
48
Appendix
Table 20: Forward-backward time on GPT-2 Large at FS=10 , TP=2 on 20 A5000 GPUs. Each optimizer has 12 runs across three allocations, with 30 timed iterations per run. The standard deviation and range width describe the 12 run means. All times are in milliseconds.
Optimizer
Operation
Reads
Frequency
Placement
AdamW
First and second moments; element-wise update
Gradient
Every step
Step
Muon
Momentum; Newton–Schulz orthogonalization
Gradient
Every step
Step
Dion
Momentum; power iteration with orthonormalization; column normalization; error feedback
Gradient
Every step
Step
Right-factor refresh in the lazy and CPU variants
Momentum
Periodic
Lagged
Shampoo, SOAP
Gradient Gram matrices
Gradient
Every step
Step
Inverse roots (Shampoo) or eigenbases (SOAP)
Gram matrices
Every f steps
Lagged
Appendix
Table 21: Dependence of optimizer operations on the current gradient, the mini-batch gradient of a weight matrix after its reduction across data-parallel ranks. Step: the operation reads the gradient and runs in the optimizer call. Fwd/bwd: it reads activations or output gradients and runs during forward-backward. Movable: it reads no gradient and can run before the optimizer call of the same iteration, once its inputs exist. Lagged: it reads statistics or momentum that include gradients, so it can leave the optimizer call only by using values from earlier steps. The first Dion row describes the standard algorithm used in our experiments. Decoupled weight decay reads only the weights and is omitted. a Our GPT-2 implementation launches the reduction behind the backward pass and runs the refresh inside the optimizer call.
Optimizer call (ms)
Pooled
Model
DP×FS×TP
GPUs
Method
Non-refresh
Refresh
Amortized
iteration
GPT-2 Large
1×10×2
20
AdamW
–
–
11.3
7764.5
Muon
–
–
623.0
8376.2
Dion r=400
–
–
1224.1
8977.3
Dobb
6.0
13.7
6.9
7760.1
GPT-2 XL
1×32×1
32
AdamW
–
–
14.5
8059.0
Appendix
Table 22: GPT-2 optimizer-call times and pooled iteration times on A5000 GPUs, in milliseconds. The GPT-2 XL runs use activation checkpointing. Refresh and Non-refresh give Dobb’s mean optimizer-call time on iterations with and without an eigenbasis refresh, over three refresh iterations per run and two at 2×16×1 . Amortized is the mean optimizer-call time, for Dobb averaged over one refresh cycle (Appendix G.1 ); the other methods perform the same update on every iteration, so their split entries are empty. Optimizer-call means cover 30 timed iterations, and 15 at 2×16×1 . Pooled iteration times are derived from the measured means (Appendix G.1 ). Dion’s label gives its rank r or rank fraction ρ .
Optimizer call (ms)
Pooled
Model
DP×FS×TP
GPUs
Method
Non-refresh
Refresh
Amortized
iteration
GPT-2 XL
1×16×1
16
AdamW
–
–
8.5
1154.3
Muon
–
–
83.7
1229.5
Dion ρ=1/4
–
–
204.6
1350.4
Dion ρ=1/16
–
–
130.0
1275.8
Dobb
5.6
22.6
7.5
1153.3
Appendix
Table 23: GPT-2 optimizer-call times and pooled iteration times on H100 GPUs, in milliseconds. Columns follow Table 22 . Optimizer-call means cover 20 timed iterations, and Dobb’s Refresh entries cover two refresh iterations per run. The measured means determine pooled iteration times through the calculation in Appendix G.1 . Dion’s label gives its rank fraction ρ . † The abstract and Section 6.2 cite this pair, which has the configuration of the GPT-2 XL loss comparison in Table 4 : Dobb with B=16 blocks per weight matrix, here one block per FSDP shard, and Dion at ρ=1/16 .
Optimizer call (ms)
FS
Method
Non-refresh
Refresh
Amortized
Pooled iteration
8
Dion
–
–
159.5
422.6
Dobb
75.7
170.0
86.2
349.3
Dobb-G- ℓ2
88.7
219.5
103.2
366.3
16
Dion
–
–
146.1
5317.7
Dobb
26.1
331.7
60.1
5231.7
Appendix
Table 24: Qwen3-8B FSDP optimizer-call times and pooled iteration times in milliseconds, with TP=1 and B=FS blocks per weight matrix. Non-refresh and Refresh give the mean optimizer-call time on iterations without and with an eigenbasis refresh; Amortized is the mean over the timed window. Dion performs the same update on every iteration, so its split entries are empty. Pooled iteration time follows Appendix G.1 , with shared components averaged over the three runs at each FS .
Optimizer call (ms)
FS
TP
Method
Non-refresh
Refresh
Amortized
Pooled iteration
32
2
Dion
–
–
1307.7
17,664.1
Dobb
32.8
1249.2
168.0
16,524.4
Dobb-G- ℓ2
53.0
1917.9
260.2
16,616.6
64
2
Dion
–
–
1287.5
18,423.8
Dobb
14.5
634.9
83.4
17,219.7
Appendix
Table 25: Qwen3-32B FSDP and TP optimizer-call times and pooled iteration times in milliseconds, with B=FS blocks per weight matrix. Columns and averaging follow Table 24 .
Distributed optimization coordinates computation across multiple agents or machines while balancing convergence, communication, privacy, and reliability.
Matrix-orthogonalization-based optimizers, exemplified by Muon, have demonstrated strong convergence behavior across a wide range of modern deep learning workloads. The matrix-aware updates offer a compelling alternative to conventional element-wise optimization, particularly as model architectures continue to grow in scale and heterogeneity. Yet contemporary distributed training infrastructure built around the assumption of element-wise optimizers is poorly matched to matrix-level optimizers such as Muon, whose updates couple entire weight matrices and require costly Newton-Schulz iterations. Vanilla Muon implementations incur more than 2x the cost of forward and backward passes. To close this gap, we present DMuon, an open-source distributed Muon implementation that integrates into existing training pipelines as a drop-in module, with no framework-level modifications. Across both embodied foundation model and large language model (LLM) training workloads, DMuon achieves a 1.48x-3.01x speedup in end-to-end step time and a 6.85x-163.00x speedup in optimizer-step time, bringing per-step latency to near-AdamW levels and enabling efficient scaling in our model training.
The Muon optimizer incurs a significant overhead cost due to its cubic-time Newton-Schulz orthogonalization step. When weights are sharded, communication overhead compounds this computational cost, eroding the benefits of Muon in many settings. We present Dion3, a revision of Muon that targets this overhead at every level of the stack. Our Gram Newton-Schulz algorithm reduces the FLOP cost of orthogonalization, our CuteDSL kernels accelerate it by exploiting symmetry, and our megabatching strategy reduces communication overhead. Moreover, we propose a simple change to the update rule that cuts costs even further: selecting only a fraction of the momentum matrix's rows to orthogonalize at each step. This update rule improves on Dion (another "compressed" version of Muon), in both speed and performance. Overall, Dion3 matches or improves on the loss achieved by Muon but reduces optimizer step time by up to 6x. Dion3 is available via the dion package (https://github.com/microsoft/dion) as a drop-in replacement for Muon.
Noah Amsel, Jack Zhang, Kwangjun Ahn +5
New York University · Microsoft Research · Princeton University +3
Shampoo is one of the leading approximate second-order optimizers: a variant of it has won the MLCommons AlgoPerf competition, and it has been shown to produce models with lower activation outliers that are easier to compress. Yet, applying Shampoo currently comes at the cost of significant computational slowdown, due to its expensive internal operations. In this paper, we take a significant step to address this shortcoming by proposing \method (for Distributed Accelerated SHampoo), a faster implementation of Distributed Shampoo based on two main new techniques: First, we show that preconditioner blocks can be stacked into 3D tensors to significantly improve GPU utilization; second, we introduce the Newton-DB iteration and the Chebyshev polynomial approximations as novel and faster approaches for computing the inverse matrix roots required by Shampoo. Along with these algorithmic contributions, we provide a first in-depth analysis of how matrix scaling critically affects Shampoo convergence. On the practical side, our GPU-aware implementation achieves up to 5.6× faster optimizer steps compared to the well-optimized Distributed Shampoo, while Newton-DB attains the lowest validation perplexity per iteration among all tested methods. Our code is available at https://github.com/IST-DASLab/DASH.
Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis +2
Institute of Science and Technology Austria · Lancaster University · Red Hat AI