🤖 machine learning

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

Het artikel introduceert DMuon, een open-source gedistribueerde implementatie van de Muon-optimizer die naadloos integreert in bestaande trainingspipelines om de optimizer-stap latentie drastisch te verminderen en een efficiëntie te bereiken die nabij die van AdamW ligt, waardoor het schaalbare gebruik van op matrix-orthogonalisatie gebaseerde optimalisatie in grote foundationmodellen mogelijk wordt.

Oorspronkelijke auteurs: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Gepubliceerd 2026-06-26

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een gigantische robotbrein (een Large Language Model of een embodied AI) traint om te leren. Om het te onderwijzen, heb je een "optimizer" nodig—een coach die naar de fouten van de robot kijkt en de interne verbindingen (gewichten) aanpast om hem slimmer te maken.

Jarenlang was de standaardcoach AdamW. Het is als een team van duizenden kleine werkers, die elk op hun beurt één klein schroefje aan de robot vastdraaien. Het is snel en efficiënt, maar het behandelt elk schroefje onafhankelijk.

Onlangs is er een nieuwe, intelligentere coach genaamd Muon ontdekt. In plaats van schroefjes één voor één te repareren, kijkt Muon naar volledige panelen van de robot (gehele matrices van data) en past deze allemaal tegelijkert aan. Deze "groepstherapie"-aanpak helpt de robot sneller te leren en een hoger prestatieniveau te bereiken. Er is echter een addertje onder het gras: Muon is extreem traag bij het draaien op een gedistribueerd systeem (een cluster van vele computers die samenwerken).

Het Problek: De "Iedereen-is-aanwezig"-vergadering-bottleneck
In een gedistribueerde trainingsopstelling wordt de hersenpan van de robot verdeeld over veel computers (GPU's).

AdamW werkt als een extern team: Computer A repareert zijn schroefjes, Computer B die van hem, en ze hoeven niet veel met elkaar te praten.
Muon werkt als een commissie: Om een paneel te repareren, moet elke computer eerst het volledige paneel te zien krijgen.

In een naïeve implementatie van Muon moet elke computer stoppen met waar hij mee bezig is, de volledige data van iedereen downloaden, de complexe wiskunde uitvoeren en vervolgens de resultaten terugsturen. Dit is als het houden van een enorme vergadering waarbij iedereen eerst een rapport van 1.000 pagina's leest voordat er een enkele beslissing kan worden genomen. Het duurt zo lang dat de "vergadering" (de optimizer stap) langer duurt dan het eigenlijke werk (de leerstappen). Sterker nog, het paper merkt op dat dit 2x meer tijd kan kosten dan het leren zelf!

De Oplossing: DMuon (Distributed Muon)
Het X Square Robot Team heeft DMuon gebouwd, een nieuw systeem dat Muon bijna net zo snel maakt als AdamW. Ze hebben dit gedaan door de manier waarop het team werkt te veranderen, niet door de wiskunde van de coach te veranderen.

Hier is hoe ze de bottleneck hebben opgelost met drie belangrijke trucs:

1. De "Gespecialiseerde Expert"-strategie (Owner-Centric Execution)

In plaats van dat iedereen probeert elk paneel te repareren, wijst DMuon één specifieke computer toe als de "Owner" (eigenaar) voor elk paneel.

Voorheen: Iedereen verzamelde de data, iedereen deed de wiskunde, iedereen verspilde tijd.
Nu: Alleen de "Owner"-computer verzamelt de data en doet de zware wiskunde. De andere computers wachten simpelweg of doen ander werk.
Analogie: Stel je een bouwplaats voor. In plaats van dat elke arbeider het hele huis probeert te bouwen, wordt er één expert toegewezen aan het dak, een andere aan de loodgieterij. De anderen proberen niet het dak te bouwen; ze leveren alleen de gereedschappen aan die de expert nodig heeft. Dit voorkomt dat iedereen twee keer hetzelfde werk doet.

2. Het "Transportband"-systeem (Communication Overlap)

Terwijl de "Owner" de wiskunde uitvoert, zijn de andere computers niet simpelweg aan het niksen. DMuon creëert een pijplijn:

Forward Pass: Terwijl de robot "denkt" (data verwerkt), stuurt het systeem op de achtergrond stilletjes de volgende set gereedschappen naar de experts.
Backward Pass: Terwijl de robot "leert van fouten", verzamelt het systeem alvast de volgende batch data voor de experts.
Analogie: Het is als een keuken in een restaurant. De chef (de Owner) is groenten aan het snijden. In plaats van te wachten tot de chef klaar is voordat de ober de volgende bestelling brengt, brengt de ober de volgende bestelling alvast terwijl de chef nog aan het snijden is. De wachttijd wordt verborgen binnen de snijtijd.

3. De "Slimme Assemblagelijn" (Batching en Tuning)

De wiskunde die Muon doet is complex. Soms zijn de stukken enorm (zoals een gigantische muur), en soms zijn ze piepklein (zoals een kleine tegel).

Het Probleem: Als je probeert een piekleine tegel op een enorme machine te verwerken, staat de machine stil. Als je een gigantische muur verwerkt, duurt het eeuwig.
De Oplossing: DMuon groepeert veel kleine tegeltjes samen in één enkele "batch" zodat de machine druk blijft. Het gebruikt ook een "tuner" die automatisch de snelste manier vindt om elke specifieke vorm van tegel te verwerken.
Analogie: In plaats van een bezorgwagen te sturen voor één enkele envelop, wacht DMuon tot hij een volle vrachtwagen aan enveloppen heeft en stuurt ze dan allemaal tegelijk. Het past ook de route van de vrachtwagen aan, afhankelijk van of de bestemming een stad of een boerderij is.

De Resultaten

Het paper heeft DMuon getest op echte modellen, waaronder een robot trainingsmodel (Wall-OSS) en een taalmodel (Qwen2.5).

Snelheid: DMuon maakte de "optimizer stap" (de vergadering) 6 tot 163 keer sneller dan de oude, naïeve methode.
Efficiëntie: De totale tijd om een stap te trainen is nu slechts 2% langzamer dan de standaard AdamW-methode.
Conclusie: DMuon stelt teams in staat om de superieure "groepstherapie"-leerstijl van Muon te gebruiken zonder de enorme tijdstraf te betalen. Het verandelt een traag, log proces in een soepel en efficiënt proces, waardoor het klaar is voor gebruik in massale AI-modellen.

Kortom, DMuon is het verkeersregelsysteem dat ervoor zorgt dat de slimste coach (Muon) op een snelweg van computers kan rijden zonder een file te veroorzaken, waardoor het eindelijk praktisch wordt om de volgende generatie AI te trainen.

Technische Samenvatting: DMuon – Efficiënte Gedistribueerde Muon-training met een Overhead Nabij Adam

1. Probleemstelling

Matrix-bewuste optimizers, in het bijzonder Muon, hebben superieure convergentie-eigenschappen aangetoond vergeleken met element-gewijze optimizers zoals AdamW door Newton–Schulz (NS) iteraties toe te passen op de momentum-geaggregeerde gradiënten van gewichtsmatrices. Deze aanpak drijft de singuliere waarden van updates naar een uniform spectrum, wat significante rekenkundige efficiëntiewinst oplevert bij het pretrainen van grote modellen (bijv. LLM's en embodied foundation models).

Het implementeren van Muon in moderne gedistribueerde trainingsomgevingen stuit echter op een kritieke systeembottleneck. Hedendaagse gedistribueerde trainingsstacks (bijv. ZeRO, FSDP) opereren onder een impliciet contract van element-gewijze optimalisatie, waarbij updates onafhankelijk worden berekend op parameter-shards. Muon schendt dit contract omdat de kern van de NS-iteratie de volledige gewichtsmatrix vereist op een enkele device om de noodzakelijke matrixproducten (specifiek $XX^\top$ ) te berekenen.

In een naïeve gedistribueerde implementatie ("Gather-then-Compute") moet elke rank:

De volledige gereduceerde gradiëntmatrix van alle shards ophalen via all-gather.
De NS-iteratie lokaal uitvoeren.
De resultaten weggooien voor de shards die hij niet bezit.

Deze aanpak brengt twee belangrijke kosten met zich mee:

Communicatie-overhead: Het reconstrueren van volledige matrices bij elke stap introduceert collectieve communicatie die schaalt met de modelgrootte en de gedistribueerde breedte.
Redundante Berekening: Elke rank voert dezelfde dure orthogonalisatie uit, wat de computationele kosten met het aantal ranks ( $D$ ) vermenigvuldigt.

In de praktijk zorgen deze kosten ervoor dat de optimizer-staptijd van gedistribueerde Muon de gecombineerde wandkloktijd van de forward- en backward-passes evenaart of zelfs overtreft, waardoor de algoritmische voordelen van Muon worden uitgehold. Dit nadeel is bijzonder ernstig voor embodied model training (bijv. VLA), waar forward-backward passes computationeel lichter zijn ten opzichte van de optimizer-stap vergeleken met standaard LLM-pretraining.

2. Methodologie: DMuon Systeemontwerp

DMuon is een open-source gedistribueerde implementatie die is ontworpen om het prestatiekloof tussen Muon en AdamW te dichten. Dit wordt bereikt door een co-designed systeemarchitectuur die communicatie-optimalisatie, load balancing en kernel-efficiëntie integreert, zonder de onderliggende trainingsframework (bijv. PyTorch/FSDP2) te wijzigen.

2.1 Owner-Centric Executiestrategie

In plaats van de berekening over alle ranks te repliceren, wijst DMuon elk matrix-parameter toe aan een enkele owner rank.

Rol van de Owner: De owner onderhoudt de gezaghebbende optimizer-state en voert de Muon-update (Newton–Schulz iteratie) precies één keer per matrix uit.
Rol van de Non-Owner: Andere ranks materialiseren de parameter alleen wanneer dat nodig is voor forward/backward berekeningen en voeren de orthogonalisatie niet uit.
Communicatieverloop:
- Forward: Owners publiceren bijgewerkte parameters asynchroon naar non-owners, gepipelined met de layer-executie.
- Backward: Gradiënten worden geaccumuleerd in gepakte buffers en gereduceerd naar de owner.
- Overlap: DMuon maakt gebruik van een twee-traps hiërarchie (intra-node en inter-node) met fijnmazige gewichtslayouts om communicatie te overlappen met berekening, waardoor collectieve congestie wordt voorkomen.

2.1 Efficiënte Gram Newton–Schulz Iteratie

Om de computationele kosten van de update aan de owner-zijde te verlagen, adopteert DMuon de Gram Newton–Schulz formulering:

Gram Space Recurrence: In plaats van te opereren in de oorspronkelijke $m \times n$ ruimte, wordt de iteratie uitgevoerd in de $m \times m$ Gram-ruimte ( $G = XX^\top$ ). Dit vermindert de complexiteit van $O(m^2n)$ naar $O(m^3)$ wanneer $m < n$ .
Benutting van Symmetrie: Omdat $G$ symmetrisch is, gebruikt DMuon een SYRK-stijl executiepad om alleen het onderste driehoeksgedeelte te berekenen, wat de rekenkundige arbeid van de dominante matrixvermenigvuldiging halveert.
Batching: Om problemen met GPU-bezetting bij kleine matrices aan te pakken, batched DMuon matrices met vergelijkbare vormen en voert de Gram NS-recurrence gezamenlijk uit, waardoor de kernel-dispatch overhead wordt geamortiseerd.

2.3 Computation-Aware Load Balancing

Omdat de kosten van de Muon-update sterk variëren met de matrixvorm, leidt een naïeve toewijzing (bijv. round-robin) tot stragglers (achterblijvers). DMuon formuleert de owner-toewijzing als een Mixed-Integer Linear Program (MILP):

Profiling: Tijdens de initialisatie benchmarkt het systeem de executietijd van verschillende batchgroottes en kernel-schema's voor elke afzonderlijke matrixvorm.
Optimalisatie: De MILP-solver minimaliseert de voorspelde makespan (maximale belasting) over de owner ranks, waardoor de zwaarste matrices worden verdeeld om de workload te balanceren.
Fallback: Voor extreem grote zoekruimtes wordt een greedische heuristiek gebruikt om de initialisatietijd te beperken.

2.4 Shape-Adaptive Kernel Autotuning

DMuon maakt gebruik van een DSL-gestuurde autotuning workflow (met behulp van TileLang en CUTE DSL) om high-performance kernels te genereren.

Het zoekt over tile-shapes, block-sizes en software pipeline-configuraties specifiek voor de doelhardware en matrixvormen.
De optimale configuratie wordt gecached en hergebruikt voor volgende stappen, wat ervoor zorgt dat het kernel-schema altijd nabij-optimaal is voor de specifieke dimensies die tijdens transformer-training voorkomen.

3. Belangrijkste Bijdragen

Fijnmazige Communicatie-optimalisatie: Een owner-centric runtime die de communicatie-overhead minimaliseert via hiërarchische scheduling en fijnmazige gewichtslayouts, waarbij de exacte optimizer-semantiek behouden blijft.
Shape-Adaptive Executie Stack: Een high-performance engine die gebatched Gram Newton–Schulz, symmetrie-bewuste kernels en autotuning combineert om heterogene matrix-workloads efficiënt te verwerken.
Computation-Aware Load Balancing: Een gemeten-kosten optimalisatiebenadering (MILP) die matrix-owners toewijst op basis van werkelijke runtime profiling in plaats van theoretische FLOP-counts.
Drop-in Integratie: Een module die met minimale inspanning integreert in bestaande trainingspipelines (bijv. FSDP2), zonder dat framework-niveau wijzigingen vereist zijn.
Productie-schaal Validatie: Succesvolle implementatie op zowel robotica (embodied) als taalmodel training workloads.

4. Experimentele Resultaten

De auteurs evalueerden DMuon op vier productie-workloads: Wall-OSS-0.5, Pi0, Wall-WM, en Qwen2.5-7B, schalend van 8 tot 256 A800 GPU's.

Optimizer Step Time: DMuon bereikt een 6.85× tot 163.00× versnelling in de optimizer-staptijd vergeleken met een naïeve "Gather-then-Compute" gedistribueerde Muon implementatie.
End-to-End Step Time: DMuon bereikt een 1.48× tot 3.01× versnelling in de totale staptijd.
Overhead vs. AdamW: De per-stap latentie van DMuon is gebracht tot nabij AdamW-niveaus, met een gemiddelde overhead van slechts +2% vergeleken met AdamW over de geëvalueerde workloads.
Scaling: Op Wall-OSS behoudt DMuon een throughput die dicht bij AdamW ligt naarmate het aantal GPU's toeneemt, terwijl naïeve Muon te maken zou krijgen met ernstige communicatiebottlenecks.

Versnelling Breakdown:

Symmetric Gram Kernel: 48% van de versnelling (halvering van de rekenkunde in dominante producten).
Owner Scheduling & Load Balancing: 32% (verwijderen van $D$ -voudige redundante berekening).
Auto-tuning & Batching: 16% (reduceren van kernel launch overhead voor kleine matrices).

5. Betekenis en Claims

Het artikel claimt dat DMuon succesvol de system mismatch tussen matrix-bewuste optimizers en gedistribueerde trainingsinfrastructuur oplost. Door redundante berekening te elimineren en communicatie te optimaliseren, stelt DMuon de convergentievoordelen van Muon (zoals hogere rekenkundige efficiëntie en betere optimalisatiegeometrie) in staat om te vertalen naar praktische wandklok-versnellingen in de training.

De auteurs positioneren DMuon niet als een nieuw algoritme, maar als een noodzakelijke systems-enabler die Muon een levensvatbare, drop-in vervanging voor AdamW maakt in grootschalige productietraining. Het werk demonstreert dat de "penalty" van matrix-bewuste optimalisatie kan worden teruggebracht tot verwaarloosbare niveaus, wat efficiënte schaling mogelijk maakt voor zowel Large Language Models als Embodied Foundation Models (VLA).

Het artikel merkt expliciet op dat DMuon de update-regel of convergentie-eigenschappen van Muon niet verandert; het adresseert strikt de systems overhead. Bij consequent is de winst het grootst in gedistribueerde settings, terwijl single-GPU training een meer bescheiden verbetering ziet (ca. 2× optimizer versnelling) vanwege het gebrek aan gedistribueerde redundantie om te elimineren.