Rohan Anil shows a tuned DistributedShampoo optimizer can match Muon's speedrun benchmarks on Modded-NanoGPT

VIEWS72.1KBOOKMARKS94LIKES229REPLIES14

I have some mixed feelings about this result:

On the one hand, it's genuinely impressive. I didn't know that Shampoo could be configured to perform this well on the benchmark.

On the other hand, the way this performance boost was achieved seems difficult to call "Vanilla," for the following reason:

According to @_arohan_, the boost depends upon fixing a numerical linear algebra issue that he observed to occur in my initial standard DistributedShampoo run. He fixed the issue by enabling the flag rank_deficient_stability_config=PseudoInverseConfig().

Here's the problem: This is an undocumented flag. It is contained within the 12,000-line DistributedShampoo codebase, but it does not appear in any user-facing documentation.

As a result, if someone tries to train a model using DistributedShampoo without either (a) knowing about this special undocumented flag or (b) being prepared to detect and fix the numerical linear algebra issues that may occur without it, then they won't be able to achieve @_arohan_'s level of Shampoo performance. This level of effort would be considered atypical for mere hyperparameter tuning.

-- [Note on Muon baseline in plot below: Rohan's post compared Shampoo to a slightly undertuned Muon baseline from 2026/05/01, which reached the target loss in 3375 steps. This resulted in a 50-step gap between Shampoo and Muon. In the figure below I'm using the up-to-date 2026/05/03 baseline, which reaches the target in 3325 steps. This results in the step-counts exactly matching between Muon and the tuned/stabilized Shampoo variant.]

rohan anil@_arohan_

Hyper-parameter tuning works really well!

Since Keller asked very nicely - I took it as a challenge to find a minimal edit distance from his config to a working config with no further code modifications, which means vanilla Meta's distributed_shampoo PyTorch package as is which won the AlgoPerf competition.

Biggest alpha here was tuning hyper-parameters apart from enabling pseudo inverse which was critical as this speed run problem produces rank deficient matrices (see those nice viz), and thanks to great help from Anna!

Hypers: lr=0.01, wd=0.1, beta2=0.9, eps=1e-15, freq=1 🙃

No Nesterov momentum was used. Oops.

All that is no longer needed - those modifications just was giving us identical training curves for same step targets because that's how math works, now the curves look different, but arrive at these end of train validation targets.

- first completed segment: step:3375/3375 val_loss:3.27656 - second completed segment: step:3375/3375 val_loss:3.27675

[1] https://github.com/facebookresearch/optimizers/issues/265#issuecomment-4668270192

Next steps for anyone looking for a late night hobby:

* I would probably fuse lot of these things, make kernels * Try it on new baselines, harder tasks: increase batch size or change architecture. * Tune frequency, change eigh to ortho items etc.

2d72.1K22994

RETWEETS5

Keller Jordan@kellerjordan0

Five recent notable Modded-NanoGPT optimization results:

Result #31: Kai Lion and Florian Hübler have improved their Muown-based run from 3075 to 2995 steps by adding NorMuon & ContraMuon modifications. 1/5

12h3.7K4512

rohan anil@_arohan_

.@kellerjordan0 posts results saying X is half way between. I educate him on optimization details, and hyper parameter tuning, and literally changes his code. I did not write any of this code.

Now he has mixed feelings and claims didn’t know Shampoo would work when he was in the same discord channel where they discussed Muon as fancy shampoo or practical shampoo in 2024. @HessianFree back me up?

Keller Jordan@kellerjordan0

I have some mixed feelings about this result:

On the one hand, it's genuinely impressive. I didn't know that Shampoo could be configured to perform this well on the benchmark.

On the other hand, the way this performance boost was achieved seems difficult to call "Vanilla," for the following reason:

According to @_arohan_, the boost depends upon fixing a numerical linear algebra issue that he observed to occur in my initial standard DistributedShampoo run. He fixed the issue by enabling the flag rank_deficient_stability_config=PseudoInverseConfig().

Here's the problem: This is an undocumented flag. It is contained within the 12,000-line DistributedShampoo codebase, but it does not appear in any user-facing documentation.

As a result, if someone tries to train a model using DistributedShampoo without either (a) knowing about this special undocumented flag or (b) being prepared to detect and fix the numerical linear algebra issues that may occur without it, then they won't be able to achieve @_arohan_'s level of Shampoo performance. This level of effort would be considered atypical for mere hyperparameter tuning.

-- [Note on Muon baseline in plot below: Rohan's post compared Shampoo to a slightly undertuned Muon baseline from 2026/05/01, which reached the target loss in 3375 steps. This resulted in a 50-step gap between Shampoo and Muon. In the figure below I'm using the up-to-date 2026/05/03 baseline, which reaches the target in 3325 steps. This results in the step-counts exactly matching between Muon and the tuned/stabilized Shampoo variant.]

2d37.7K15685

You Jiacheng@YouJiacheng

When I discussed with @Jianlin_S , he said computing P^{-1/2} separately is unstable (context: Gram-NS; and that's why it needs restart) and I fully agree.

AND, he proposed an iterative method to compute P^{-1/2}G stably in https://kexue.fm/archives/11158.

Keller Jordan@kellerjordan0

I have some mixed feelings about this result:

On the one hand, it's genuinely impressive. I didn't know that Shampoo could be configured to perform this well on the benchmark.

On the other hand, the way this performance boost was achieved seems difficult to call "Vanilla," for the following reason:

According to @_arohan_, the boost depends upon fixing a numerical linear algebra issue that he observed to occur in my initial standard DistributedShampoo run. He fixed the issue by enabling the flag rank_deficient_stability_config=PseudoInverseConfig().

Here's the problem: This is an undocumented flag. It is contained within the 12,000-line DistributedShampoo codebase, but it does not appear in any user-facing documentation.

As a result, if someone tries to train a model using DistributedShampoo without either (a) knowing about this special undocumented flag or (b) being prepared to detect and fix the numerical linear algebra issues that may occur without it, then they won't be able to achieve @_arohan_'s level of Shampoo performance. This level of effort would be considered atypical for mere hyperparameter tuning.

-- [Note on Muon baseline in plot below: Rohan's post compared Shampoo to a slightly undertuned Muon baseline from 2026/05/01, which reached the target loss in 3375 steps. This resulted in a 50-step gap between Shampoo and Muon. In the figure below I'm using the up-to-date 2026/05/03 baseline, which reaches the target in 3325 steps. This results in the step-counts exactly matching between Muon and the tuned/stabilized Shampoo variant.]

2d11.2K6753

rohan anil@_arohan_

Newton Schulz is the fast operator for Matrix Sign the special case when b2=0.0 - this makes Muon fast. This is very cool!

Denman-Beavers coupled newton iteration for inverse square roots and fourth roots is what you need when b2>0.0 - this is what you need to make it really fast.

It’s available in the pytorch package.

James MMatrix@JamesWhate89993

Fun fact: Jeremy’s first Muon research thread (https://x.com/jxbz/status/1819846348130418706 ) back in Aug 2024 exactly started from Shampoo winning AlgoPerf, then gave the key observation: “Shampoo preconditioning exactly orthogonalizes the gradient matrices.”

Then the Old Optimizer, New Norm paper derived Newton–Schulz, the key innovation of Muon, as a “computational strategy for Shampoo”.

So “Muon is not exactly Shampoo” is a dodge. Of course it isn’t byte-for-byte Shampoo. That was never the point.

The point is that Muon was born downstream of Shampoo, as a practical Shampoo-line variant, then got hype-maxxed into “the first big optimizer breakthrough after Adam.”

No. Shampoo was the breakthrough. The Shampoo line led to amazing optimizers like SOAP, which was invented before Muon, and btw still leads the current ModdedNanoGPT optimizer track record

2d10.5K6946

Susan Zhang@suchenzang

we should really be thanking arabic numerals for making zeros happen

(screenshot posted with permission, and brings the thread full circle :))

Aurko Roy@aurko79

As someone who watched Rohan train LLMs with shampoo and brainstorm ideas for making it cheaper/more scalable back in 2021-2022 (low-rank shampoo for very large weights etc), I think it's fair to say that Shampoo was certainly a step change over Adam.

That said Muon also introduced several innovations and it may be an over-simplification to say that it's just Shampoo with b2=0.

- For one the use of Nesterov momentum - Second the observation that computing inverse roots is inherently more numerically unstable than orthogonalization - Clever use of Newton-Schulz which makes it more GPU friendly than eigh or QR or coupled newton - Scaling the update rms norm to match 0.2 (Adam update rms in the limit) which made grafting with Adam unnecessary

I think both optimizers are awesome and I learned a lot from studying both of them. Progress doesn't happen in a vacuum, we all stand on the shoulders of giants.

2d19.2K11532

Dimitris Papailiopoulos@DimitrisPapail

Oh là là, SGD may get there. Fumble 5 may in fact be good

Dimitris Papailiopoulos@DimitrisPapail

I have a silly question. How far off is vanilla SGD with well tuned learning rate schedule and batch size? I'd love to see wall clock on the x axis

1d22.8K6627

Dimitris Papailiopoulos@DimitrisPapail

I have a silly question. How far off is vanilla SGD with well tuned learning rate schedule and batch size? I'd love to see wall clock on the x axis

Keller Jordan@kellerjordan0

I have some mixed feelings about this result:

On the one hand, it's genuinely impressive. I didn't know that Shampoo could be configured to perform this well on the benchmark.

On the other hand, the way this performance boost was achieved seems difficult to call "Vanilla," for the following reason:

According to @_arohan_, the boost depends upon fixing a numerical linear algebra issue that he observed to occur in my initial standard DistributedShampoo run. He fixed the issue by enabling the flag rank_deficient_stability_config=PseudoInverseConfig().

Here's the problem: This is an undocumented flag. It is contained within the 12,000-line DistributedShampoo codebase, but it does not appear in any user-facing documentation.

As a result, if someone tries to train a model using DistributedShampoo without either (a) knowing about this special undocumented flag or (b) being prepared to detect and fix the numerical linear algebra issues that may occur without it, then they won't be able to achieve @_arohan_'s level of Shampoo performance. This level of effort would be considered atypical for mere hyperparameter tuning.

-- [Note on Muon baseline in plot below: Rohan's post compared Shampoo to a slightly undertuned Muon baseline from 2026/05/01, which reached the target loss in 3375 steps. This resulted in a 50-step gap between Shampoo and Muon. In the figure below I'm using the up-to-date 2026/05/03 baseline, which reaches the target in 3325 steps. This results in the step-counts exactly matching between Muon and the tuned/stabilized Shampoo variant.]

1d30.6K7025

Dimitris Papailiopoulos@DimitrisPapail

taught grandpa SGD to speedrun GPT tonight 2x slower than Muon.

We are 0.10 nats from god and $300 of runpod credits poorer... yet doing the single most dangerous thing I can:

leaving the agents (codex and CC with fumble) run overnight 🫡

see wassup tomorrow. either 3.28 at 550s or go home with 300$ less

Dimitris Papailiopoulos@DimitrisPapail

We may be only 250$ (aka 10 hours) away from the machine god uber optimizing SGD to close the gap

1d10.9K5323

You Jiacheng@YouJiacheng

Muon is Shampoo (2018)'s one-sided & β1=β2 variant without variance-damping. Muon is Stochastic Spectral Descent (2015) with Momentum. one-sided Shampoo vs. Muon = Adam vs. SignSGD. SSD vs. Muon = SGD vs. SGDM.

Note: but Muon is not Shampoo w/ b2=0.

2d5.6K5923

Tim Lau@timlautk

I have an alternative view based on the principle of symmetry and bi-orthogonal equivariance (Sections 3.2 & 3.8 of https://arxiv.org/abs/2605.18106). Not an attempt to attack Shampoo but just trying to understand the underlying geometric behavior of different optimizers.

Muon itself is an example of spectral optimizers such that whenever a matrix gradient or momentum direction is represented in a rotated basis (PGQᵀ or PMQᵀ for orthogonal P, Q), their update are transformed in the same way. Other instances also include Scion, PolarGrad and even vanilla SGD.

In contrast, in Section 3.8.1, it is shown that Shampoo (b2 != 0) is generally not a spectral operator of the current gradient/momentum alone, as the eigenvectors of L and R are not the same as the left and right singular vectors of G (or M). Since G (or M) is the (smoothed) steepest descent direction, any update map acts on it should preserve its singular vectors (i.e., a spectral operator). Ideally, the update map should preserve the update directions but potentially change the update magnitude of the directions, interpreted as some form of preconditioning.

Of course, there is gradient noise but I'd treat the M = EMA(G) as a denoiser rather than EMA(GGᵀ) or EMA(GᵀG). I think the so-called variance adaptation/damping (also in Adam) should rather be understood as an approximation of preconditioning, but such variance adaptation/damping isn't bi-orthogonal equivariant unless b2 = 0. Likewise, the elementwise variance-adaptation map in Adam is also not bi-orthogonal equivariant.

Following this argument, I would rather say that it boils down to determine the "best" spectral optimizer, i.e., the "best" absolutely symmetric map of the singular values of G (or M). Setting all singular values to one like in Muon might not be the best. I suspect that the "best" spectral optimizer might not exist and could be layer- or architecture-dependent, but there are some recent works on trying different powers of the singular values.

This is only my personal take and hope it would be constructive for developing a better understanding of the underlying behavior of different existing optimizers.

1d3.3K2824

kache@yacineMTB

Spending 100x the compute of training to find hparams that make the training .5% more efficient

Keller Jordan@kellerjordan0

I have some mixed feelings about this result:

On the one hand, it's genuinely impressive. I didn't know that Shampoo could be configured to perform this well on the benchmark.

On the other hand, the way this performance boost was achieved seems difficult to call "Vanilla," for the following reason:

According to @_arohan_, the boost depends upon fixing a numerical linear algebra issue that he observed to occur in my initial standard DistributedShampoo run. He fixed the issue by enabling the flag rank_deficient_stability_config=PseudoInverseConfig().

Here's the problem: This is an undocumented flag. It is contained within the 12,000-line DistributedShampoo codebase, but it does not appear in any user-facing documentation.

As a result, if someone tries to train a model using DistributedShampoo without either (a) knowing about this special undocumented flag or (b) being prepared to detect and fix the numerical linear algebra issues that may occur without it, then they won't be able to achieve @_arohan_'s level of Shampoo performance. This level of effort would be considered atypical for mere hyperparameter tuning.

-- [Note on Muon baseline in plot below: Rohan's post compared Shampoo to a slightly undertuned Muon baseline from 2026/05/01, which reached the target loss in 3375 steps. This resulted in a 50-step gap between Shampoo and Muon. In the figure below I'm using the up-to-date 2026/05/03 baseline, which reaches the target in 3325 steps. This results in the step-counts exactly matching between Muon and the tuned/stabilized Shampoo variant.]

1d15.3K8317

Dimitris Papailiopoulos@DimitrisPapail

BOOM shakalaka

now shortening the timeline :D

Dimitris Papailiopoulos@DimitrisPapail

Serious Greek Drama.

This is what hyperparam overfitting looks like

12h13.6K7416

Dimitris Papailiopoulos@DimitrisPapail

We may be only 250$ (aka 10 hours) away from the machine god uber optimizing SGD to close the gap

Dimitris Papailiopoulos@DimitrisPapail

Oh là là, SGD may get there. Fumble 5 may in fact be good

1d29.4K4719

Dimitris Papailiopoulos@DimitrisPapail

It is looking very very good for Stochastic Grandpa

Dimitris Papailiopoulos@DimitrisPapail

We may be only 250$ (aka 10 hours) away from the machine god uber optimizing SGD to close the gap

14h14.9K6211

Jeremy Cohen@deepcohen

@DimitrisPapail You might find this interesting: https://arxiv.org/abs/2407.07972

Dimitris Papailiopoulos@DimitrisPapail

Oh là là, SGD may get there. Fumble 5 may in fact be good

1d1.5K1122

Keller Jordan@kellerjordan0

I created this new speedrun track, which compares results in terms of steps rather than wallclock, specifically to give a fair chance to optimizers other than Muon.

Happy to see the resulting accumulation of public knowledge!

Keller Jordan@kellerjordan0

Result #35: @_arohan_ has achieved a significant >700-step improvement to the best known DistributedShampoo config.

This was achieved by switching to one-sided Shampoo, numerically stabilizing the run using an undocumented pseudoinverse flag, and retuning other hparams. 4/5

12h8.1K7212

rohan anil@_arohan_

“Shampoo falls halfway between Muon & Adam” was the very confident claim.

This has been addressed. It had to be done.

Nesterov largely have not been useful in language models from my experience and now for that I got the answer as well. Not enough effort has gone to these submissions.

This is why we do tuning budgets anf made as fair comparison as possible @zacharynado

https://arxiv.org/abs/2502.15015

2d1.4K2818

You Jiacheng@YouJiacheng

😂lol now Shampoo is again in the middle between Muon and Adam. next round, who will tune Shampoo to match / surpass Muon?

Konstantin Mishchenko@konstmish

I just submitted a PR to modded-nanogpt with better hyperparams. With them, Muon can reach the target loss after 3250 steps instead of 3325. Always tune your baseline well when doing research. Weak baselines can make any idea look promising

1d7.1K488

Keller Jordan@kellerjordan0

Result #33: @varunneal has provided a PSGD-Kron (Pooladzandi & Li 2024; Li 2024; Li 2022; Li 2018; Li 2015) run! It uses hyperball optimization. 3/5

Keller Jordan@kellerjordan0

Result #32: @mihai673 has achieved a 30-step improvement over the old 2026/05/09 record by adding a SODA (Pethick et al. 2026)-style anchor towards init.

It is unknown whether this technique can also improve the current record. 2/5

12h2.2K475