noahamsel retweeted
TLDR: we integrate many computational and statistical improvements for muon-style updates here: github.com/microsoft/dion .
Comp: Gram Newton Schulz, Row Selection, Megabatching, Cudagraph Capture, Symmetric Kernels
Stat: Joint learning rate+fraction, row norm, fused rounding
Newton-Schulz is what makes Muon work, but it's also a massive pain compared to Adam. *Dion3* makes it ~6x faster and more parallelizable too.
Combines our work on row-filtered gradients, Gram-NS alg, and GPU kernels in one package. Try pip install dion!
arxiv.org/abs/2608.11612
Newton-Schulz is what makes Muon work, but it's also a massive pain compared to Adam. *Dion3* makes it ~6x faster and more parallelizable too.
Combines our work on row-filtered gradients, Gram-NS alg, and GPU kernels in one package. Try pip install dion!
arxiv.org/abs/2608.11612
Getting improvements across the whole stack was only possible thanks to my expert collaborators and mentors at Microsoft Research, Princeton, and beyond #synergy #teamwork #distributedorthonormalization
@Ali_NT99 @jcz42 Austin Feng @KwangjunA @_berlinchen @tri_dao @JohnCLangford
noahamsel retweeted
PR for SM120:
github.com/Dao-AILab/quack/p…
Symmetric Gemm: github.com/Dao-AILab/quack/b…
Gram-Newton Schulz:
github.com/Dao-AILab/gram-ne…
Research Blog:
dao-lab.ai/blog/2026/gram-ne…
noahamsel retweeted
It's my favorite kind of work: linear algebra insight + fast kernels.
When playing w Muon a while ago, we were thinking why not speed it up by operating on the small square matrix X X^T instead of the large rectangular matrix X. Jack, Noah, and Berlin spent many months understanding eigenvalues/vectors of the intermediate matrices in Muon, and finally came up with a simple and elegant algo to make this work.
We made Muon run up to 2x faster for free!
Introducing Gram Newton-Schulz: a mathematically equivalent but computationally faster Newton-Schulz algorithm for polar decomposition.
Gram Newton-Schulz rewrites Newton-Schulz such that instead of iterating on the expensive rectangular X matrix, we iterate on the small, square, symmetric XX^T Gram matrix to reduce FLOPs. This allows us to make more use of fast symmetric GEMM kernels on Hopper and Blackwell, halving the FLOPs of each of those GEMMs.
Gram Newton-Schulz is a drop-in replacement of Newton-Schulz for your Muon use case: we see validation perplexity preserved within 0.01, and share our (long!) journey stabilizing this algorithm and ensuring that training quality is preserved above all else.
This was a super fun project with @noahamsel, @berlinchen, and @tri_dao that spanned theory, numerical analysis, and ML systems! Blog and codebase linked below 🧵
Announcing Gram Newton-Schulz, a new way to implement Muon that's 2x faster
Trick 1: rejigger Newton-Schulz to replace rectangular matmuls with square symmetric ones
Trick 2: collaborate with CUDA kings @jcz42 and @_berlinchen from @tri_dao's lab
Blog: dao-ailab.github.io/blog/202…
We made Muon run up to 2x faster for free!
Introducing Gram Newton-Schulz: a mathematically equivalent but computationally faster Newton-Schulz algorithm for polar decomposition.
Gram Newton-Schulz rewrites Newton-Schulz such that instead of iterating on the expensive rectangular X matrix, we iterate on the small, square, symmetric XX^T Gram matrix to reduce FLOPs. This allows us to make more use of fast symmetric GEMM kernels on Hopper and Blackwell, halving the FLOPs of each of those GEMMs.
Gram Newton-Schulz is a drop-in replacement of Newton-Schulz for your Muon use case: we see validation perplexity preserved within 0.01, and share our (long!) journey stabilizing this algorithm and ensuring that training quality is preserved above all else.
This was a super fun project with @noahamsel, @berlinchen, and @tri_dao that spanned theory, numerical analysis, and ML systems! Blog and codebase linked below 🧵
How can classical numerical analysis help train deep nets faster? Climb aboard the Polar Express to find out... arxiv.org/abs/2505.16932
joint with @davpersson @gowerrobert + Chris Musco
