Muon

Basic formulation

• $G_t=\nabla {\mathcal{L}}_t(W_{t-1})$
• $M_t=\mu M_{t-1}+G_t$
• $O_t=NS(M_t)$
• ${\theta }_t\leftarrow {\theta }_{t-1}-\eta O_t$

Extensions

• MuonBP: Faster Muon via Block-Periodic Orthogonalization

  • Using Muon on large sharded models creates extra communication overhead from gather/scatter operations on sharded matrices. Turns out you can fix this by doing full Muon updates periodically (but not skipping them entirely) and using local Muon computation the rest of the time.

• https://www.essential.ai/blog/infra layer sharding for large scale training with muon

• Various approaches to parallelizing Muon (mainhorse blog)

• Practical Efficiency of Muon for Pretraining

  • essentialAI infra tips for muon