Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases

Shi-Yuk Fong; Shuhua Yu; Jiawei Zhang

cs.LGOct 29, 2026

Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases

Authors: Shi-Yuk Fong, Shuhua Yu, Jiawei Zhang

Organizations: Department of Computer Sciences, University of Wisconsin–Madison, Madison, WI, USA · Meta FAIR, Menlo Park, CA, USA

Abstract

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

\rho

. For a weight matrix of size

m\times 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\rho

. Payloads are derived for FSDP with

\mathrm{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 ).

Panel of Figure 2 — 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(-\sum_{i}p_{i}\log p_{i})$ as a percentage of the factor dimension, $n$ for ${\bm{A}}_{t}$ and $m$ for ${\bm{G}}_{t}$ , where $p_{i}$ 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 ${\bm{Y}}$	${\bm{Y}}$	$\operatorname{sign}({\bm{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\times 768

gradients. (a) A matrix

{\bm{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

{\bm{D}}{\bm{X}}^{\top}

whose output-gradient second moment

{\bm{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

{\bm{Y}}

in (a). Orthogonality is

\|{\bm{O}}\|_{*}/(\sqrt{\min(m,n)}\,\|{\bm{O}}\|_{F})

for a nonzero

m\times n

matrix

{\bm{O}}

, where

\|\cdot\|_{*}

and

\|\cdot\|_{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	$\mathcal{O}(mn)$
Dion, rank $r$	$\mathcal{O}(r(m+n))$
Dobb	$0$

Table 3: Leading optimizer-specific cross-shard payload for one

m\times 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,

n^{2}/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	$\mathrm{DP}$	$\mathrm{FS}$	$\mathrm{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

\mathrm{DP}

data-parallel replicas and FSDP shard degree

\mathrm{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,

(\text{Refresh}+8\cdot\text{Non-refresh})/9

. Iteration is the pooled iteration time defined in Appendix G.1 . Dion uses rank

r=400

on Large; rank fraction

\rho=1/3

on XL (A5000), 6B, and 8B;

\rho=1/4

on XL (H100) and 3B; and on Qwen3, a rank fraction matched to the block width

1024/\mathrm{FS}

(Appendix G.5 ). For GPT-2 XL on H100, Table 23 also reports Dion at

\rho=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
${\bm{W}}_{t}$ , $t$	Weight matrix of a layer ${\bm{y}}={\bm{W}}{\bm{x}}$ at optimizer step $t$	Section 3
$\mathcal{B}_{t}$ , $N$	Mini-batch at step $t$ , and its number of examples or token positions	Section 3
$f_{\mathcal{B}_{t}}$	Mean loss on $\mathcal{B}_{t}$	Equation 1
${\bm{x}}_{i}$ , ${\bm{g}}_{i}$	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- $\ell_{2}$	$\mathcal{O}(n)$	$0$
Dobb-G- $\ell_{2}$	$0$	$\mathcal{O}(m)$

Appendix

Table 7: Leading optimizer-specific cross-shard payload for the basis and normalization ablations on one

m\times n

matrix, with each block within one shard. The

\ell_{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\times 10^{-3}$	$3.54\times 10^{-3}$	$2.5\times 10^{-3}$	$2.5\times 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	$\\|{\bm{O}}\\|_{F}^{2}$	$s$	RMS of $s{\bm{O}}$	GPT-2
Full-rank orthogonalized Muon	$\min(m,n)$	$\tau\sqrt{\max(m,n)}$	$\tau$	0.2
Dion, $r$ orthonormal left / unit-norm right columns	$r$	$\tau\sqrt{\max(m,n)}$	$\tau\sqrt{r/\min(m,n)}$	$0.2\sqrt{\rho}$
RMNP, unit-norm rows	$m$	$\tau\sqrt{\max(m,n)}$	$\tau\sqrt{\max(m,n)/n}$	0.2 or 0.4 a
Dobb, Dobb- $\ell_{2}$	$n$	$\tau\sqrt{\max(m,n)}$	$\tau\sqrt{\max(m,n)/m}$	0.2 or 0.4 b
Dobb-G, Dobb-G- $\ell_{2}$	$m$	$\tau\sqrt{\max(m,n)}$	$\tau\sqrt{\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

\tau=0.2

and the rank fraction

\rho=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

\rho<1

$N_{\mathrm{param}}$	Layers	Width	Heads	Steps	$D$	Grid $j$	Runs	$\widehat{\gamma}^{\star}$ ( $10^{-3}$ )
Swept sizes
62M	6	640	10	5,000	2.46B	$-4\ldots 4$	$9+5^{a}$	4.16 [3.74, 4.45]
124M	12	768	12	10,000	4.92B	$-4\ldots 4$	9	3.56 [3.42, 3.72]
219M	18	896	14	15,000	7.37B	$-3\ldots 1$	5	3.56 [3.31, 3.83]
354M	24	1024	16	20,000	9.83B	$-3\ldots 1$	5	3.13 [2.43, 3.60]
Extrapolated sizes

Appendix

Table 10: Model sizes and learning-rate sweeps.

N_{\mathrm{param}}

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,

\widehat{\gamma}^{\star}

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,

\widehat{\gamma}^{\star}

is the extrapolation of Equation 11 with a 68% residual-bootstrap interval; both models are trained at

\gamma=2.5\times 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

\pm 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,

\gamma=2.5\times 10^{-3}

(

j=-2

). The shaded region lies beyond the largest swept size.

	Small			Medium			Large			XL
$B$	Dobb	Dobb- $\ell_{2}$	$\Delta$	Dobb	Dobb- $\ell_{2}$	$\Delta$	Dobb	Dobb- $\ell_{2}$	$\Delta$	Dobb	Dobb- $\ell_{2}$	$\Delta$
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-

\ell_{2}

(column

\ell_{2}

). The Dobb columns repeat Table 4 ;

\Delta

is the Dobb-

\ell_{2}

loss minus the Dobb loss.

	Activation basis		Output-gradient basis
$B$	Dobb	Dobb- $\ell_{2}$	Dobb-G	Dobb-G- $\ell_{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
$\ell_{\infty}$	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
$\ell_{2}$	Activation second moment (Dobb- $\ell_{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	$8n^{2}/B$
fp32	bf16 (principal)	3.017	360.4	$6n^{2}/B$
bf16	fp32	3.028	361.0	$6n^{2}/B$
bf16	bf16	3.027	360.9	$4n^{2}/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

\mathrm{FS}=10

\mathrm{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	$\mathrm{DP}\times\mathrm{FS}\times\mathrm{TP}$	GPUs	Method	Non-refresh	Refresh	Amortized	iteration
GPT-2 Large	$1\times 10\times 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\times 32\times 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\times 16\times 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\times 16\times 1

. Pooled iteration times are derived from the measured means (Appendix G.1 ). Dion’s label gives its rank

r

or rank fraction

\rho

				Optimizer call (ms)			Pooled
Model	$\mathrm{DP}\times\mathrm{FS}\times\mathrm{TP}$	GPUs	Method	Non-refresh	Refresh	Amortized	iteration
GPT-2 XL	$1\times 16\times 1$	16	AdamW	–	–	8.5	1154.3
			Muon	–	–	83.7	1229.5
			Dion $\rho=1/4$	–	–	204.6	1350.4
			Dion $\rho=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

\rho

. † 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

\rho=1/16

$\mathrm{FS}$	Method	Non-refresh	Refresh	Amortized	Pooled iteration
		Optimizer call (ms)
8	Dion	–	–	159.5	422.6
	Dobb	75.7	170.0	86.2	349.3
	Dobb-G- $\ell_{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

\mathrm{TP}=1

and

B=\mathrm{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

\mathrm{FS}

$\mathrm{FS}$	$\mathrm{TP}$	Method	Non-refresh	Refresh	Amortized	Pooled iteration
			Optimizer call (ms)
32	2	Dion	–	–	1307.7	17,664.1
		Dobb	32.8	1249.2	168.0	16,524.4
		Dobb-G- $\ell_{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=\mathrm{FS}

blocks per weight matrix. Columns and averaging follow Table 24 .

Distributed Optimization

Distributed optimization coordinates computation across multiple agents or machines while balancing convergence, communication, privacy, and reliability.

See all 93 papers

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead
Jun 25, 2026Vincent Chen, Starrick Liu, Regis Cheng +8Deep Learning Optimization Distributed Optimization
X SQUARE ROBOT TEAM
Dion3: Full-Stack Orthogonal Updates
Aug 12, 2026Noah Amsel, Jack Zhang, Kwangjun Ahn +5Newton-Schulz Iteration GPU Kernel Optimization
New York University · Microsoft Research · Princeton University +3
DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers
Feb 2, 2026Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis +2Deep Learning Optimization GPU Acceleration
Institute of Science and Technology Austria · Lancaster University · Red Hat AI

Jun 25, 2026J/K move · Enter open · S save

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang+7

X SQUARE ROBOT TEAM

Deep Learning Optimization Distributed Optimization Muon Optimizer Distributed Training

Open paper AskPDF arXiv

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.

Abstract

Distributed Optimization

Explore similar work

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

Dion3: Full-Stack Orthogonal Updates

DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers