🤖 machine learning

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

Das Paper stellt DMuon vor, eine Open-Source-distribuierte Implementierung des Muon-Optimierers, der sich nahtlos in bestehende Trainings-Pipelines integriert, um die Optimizer-Step-Latenz drastisch zu reduzieren und eine nahezu AdamW-ähnliche Effizienz zu erreichen, wodurch dadurch die skalierbare Anwendung der auf Matrix-Orthogonalisierung basierenden Optimierung in großen Foundation-Modellen ermöglicht wird.

Ursprüngliche Autoren: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Veröffentlicht 2026-06-26

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie trainieren das Gehirn eines riesigen Roboters (ein Large Language Model oder eine eingebettete KI), um zu lernen. Um es zu lehren, benötigen Sie einen „Optimierer“ – einen Coach, der auf die Fehler des Roboters schaut und seine internen Verbindungen (Gewichte) anpasst, um ihn intelligenter zu machen.

Jahrelang war der Standard-Coach AdamW. Er ist wie ein Team aus tausenden winzigen Arbeitern, von denen jeder jeweils eine winzige Schraube am Roboter repariert. Er ist schnell und effizient, aber er behandelt jede Schraube unabhängig voneinander.

Vor kurzem wurde ein neuer, klügerer Coach namens Muon entdeckt. Anstatt Schrauben einzeln zu reparieren, betrachtet Muon ganze Paneele des Roboters (ganze Matrizen von Daten) und passt sie alle gemeinsam an. Dieser „Gruppentherapie“-Ansatz hilft dem Roboter, schneller zu lernen und eine höhere Leistungsfähigkeit zu erreichen. Es gibt jedoch einen Haken: Muon ist unglaublich langsam, wenn es auf einem verteilten System (einem Cluster aus vielen Computern, die zusammenarbeiten) läuft.

Das Problem: Der Flaschenhals der „Alle-müssen-teilnehmen“-Sitzung
In einem verteilten Trainingssetup ist das Gehirn des Roboters auf viele Computer (GPUs) aufgeteilt.

AdamW arbeitet wie ein Remote-Team: Computer A repariert seine Schrauben, Computer B seine, und sie müssen nicht viel miteinander reden.
Muon arbeitet wie ein Komitee: Um ein Paneel zu reparieren, muss jeder Computer zuerst das gesamte Paneel sehen.

In einer naiven Implementierung von Muon muss jeder einzelne Computer aufhören, was er gerade tut, die vollständigen Daten von allen anderen herunterladen, die komplexe Mathematik durchführen und dann die Ergebnisse zurücksenden. Das ist wie eine massive Sitzung, bei der jeder einen 1.000-seitigen Bericht liest, bevor eine einzige Entscheidung getroffen werden kann. Es dauert so lange, dass die „Sitzung“ (der Optimierungsschritt) länger dauert als die eigentliche Arbeit (die Lernschritte). Tatsächlich stellt das Paper fest, dass dies 2-mal mehr Zeit kosten kann als das eigentliche Lernen!

Die Lösung: DMuon (Distributed Muon)
Das X Square Robot Team hat DMuon entwickelt, ein neues System, das Muon fast so schnell wie AdamW laufen lässt. Sie haben dies erreicht, indem sie die Art und Weise geändert haben, wie das Team arbeitet, nicht durch eine Änderung der Mathematik des Coaches.

Hier ist, wie sie den Flaschenhals mit drei Haupttricks gelöst haben:

1. Die Strategie des „beauftragten Experten“ (Besitzerzentrierte Ausführung)

Anstatt dass alle versuchen, jedes Paneel zu reparieren, weist DMuon jedem Paneel einen spezifischen Computer als „Besitzer“ (Owner) zu.

Vorher: Alle sammelten die Daten, alle machten die Mathematik, alle verschwendeten Zeit.
Jetzt: Nur der „Besitzer“-Computer sammelt die Daten und erledigt die schwere Mathematik. Die anderen Computer warten einfach ab oder leisten andere Arbeit.
Analogie: Stellen Sie sich eine Baustelle vor. Anstatt dass jeder Arbeiter versucht, das ganze Haus zu bauen, wird einem Experten das Dach und einem anderen die Sanitäranlagen zugewiesen. Die anderen versuchen nicht, das Dach zu bauen; sie reichen dem Experten nur die Werkzeuge, die er braucht. Dies verhindert, dass alle die gleiche Arbeit doppelt ausführen.

2. Das „Förderband“-System (Kommunikations-Überlappung)

Während der „Besitzer“ die Mathematik betreibt, sind die anderen Computer nicht einfach nur untätig. DMuon erstellt eine Pipeline:

Forward Pass: Während der Roboter „denkt“ (Daten verarbeitet), sendet das System im Hintergrund bereits den nächsten Satz an Werkzeugen an die Experten.
Backward Pass: Während der Roboter „aus Fehlern lernt“, sammelt das System bereits die nächste Charge an Daten für die Experten ein.
Analogie: Es ist wie in einer Restaurantküche. Der Koch (der Besitzer) schneidet Gemüse. Anstatt darauf zu warten, dass der Koch fertig ist, bevor der Kellner die nächste Bestellung bringt, bringt der Keller die nächste Bestellung, während der Koch noch am Schneiden ist. Die Wartezeit wird in die Zeit des Schneidens „versteckt“.

3. Die „Intelligente Montagestraße“ (Batching und Tuning)

Die Mathematik, die Muon betreibt, ist komplex. Manchmal sind die Stücke riesig (wie eine gigantische Wand), und manchmal sind sie winzig (wie eine kleine Kachel).

Das Problem: Wenn man versucht, eine winzige Kachel auf einer riesigen Maschine zu verarbeiten, steht die Maschine still. Wenn man eine riesige Wand verarbeitet, dauert es ewig.
Die Lösung: DMuon gruppiert viele winzige Kacheln zu einem einzigen „Batch“ zusammen, damit die Maschine beschäftigt bleibt. Zudem verwendet es einen „Tuner“, der automatisch den schnellsten Weg findet, um jede spezifische Form der Kachel zu bearbeiten.
Analogie: Anstatt für einen einzelnen Umschlag einen Lieferwagen zu schicken, wartet DMuon, bis es eine volle Lkw-Ladung an Umschlägen hat, und schickt sie dann alle auf einmal. Es passt auch die Route des LKWs an, je nachdem, ob das Ziel eine Stadt oder ein Bauernhof ist.

Die Ergebnisse

Das Paper hat DMuon an realen Modellen getestet, einschließlich eines Robotertrainingsmodells (Wall-OSS) und eines Sprachmodells (Qwen2.5).

Geschwindigkeit: DMuon machte den „Optimierungsschritt“ (die Sitzung) 6- bis 163-mal schneller als die alte, naive Methode.
Effizienz: Die Gesamtzeit für einen Trainingsschritt ist nun nur noch 2 % langsamer als die Standard-AdamW-Methode.
Fazit: DMuon ermöglicht es Teams, den überlegenen „Gruppentherapie“-Lernstil von Muon zu nutzen, ohne den massiven Zeitverlust in Kauf nehmen zu müssen. Es verwandelt einen langsamen, klobigen Prozess in einen reibungslosen, effizienten Ablauf und macht ihn somit bereit für den produktiven Einsatz in massiven KI-Modellen.

Kurz gesagt: DMuon ist das Verkehrsleitsystem, das es dem klügsten Coach (Muon) ermöglicht, auf einer Autobahn aus Computern zu fahren, ohne einen Stau zu verursachen, und macht ihn so endlich praktikabel für das Training der nächsten Generation von KI.

Technisches Resümee: DMuon – Effizientes verteiltes Muon-Training mit nahezu Adam-Overhead

1. Problemstellung

Matrix-bewusste Optimierer, insbesondere Muon, haben im Vergleich zu elementweisen Optimierern wie AdamW überlegene Konvergenzeigenschaften demonstriert, indem sie Newton–Schulz (NS)-Iterationen auf die momentum-aggregierten Gradienten von Gewichtmatrizen anwenden. Dieser Ansatz treibt die Singularwerte der Updates hin zu einem uniformen Spektrum und bietet signifikante Vorteile bei der Recheneffizienz beim Pretraining großer Modelle (z. B. LLMs und embodied Foundation Models).

Die Bereitstellung von Muon in modernen verteilten Trainingsumgebungen stößt jedoch auf einen kritischen System-Engpass. Zeitgenössische verteilte Trainingsstacks (z. B. ZeRO, FSDP) operieren unter einem impliziten Vertrag der elementweisen Optimierung, bei dem Updates auf den Parameter-Shards unabhängig berechnet werden. Muon verletzt diesen Vertrag, da seine Kern-NS-Iteration die vollständige Gewichtstrix auf einem einzelnen Gerät benötigt, um die notwendigen Matrixprodukte (speziell $XX^\top$ ) zu berechnen.

In einer naiven verteilten Implementierung („Gather-then-Compute“) muss jeder Rank:

Die vollständige reduzierte Gradientenmatrix aus allen Shards mittels All-gather sammeln.
Die NS-Iteration lokal durchführen.
Die Ergebnisse für die Shards verwerfen, die ihm nicht gehören.

Dies verursacht zwei wesentliche Kosten:

Kommunikations-Overhead: Die Rekonstruktion vollständiger Matrizen in jedem Schritt führt zu kollektiver Kommunikation, die mit der Modellgröße und der Verteilungsbreite skaliert.
Redundante Berechnung: Jeder Rank führt dieselbe teure Orthogonalisierung durch, was die Rechenkosten um die Anzahl der Ranks ( $D$ ) multipliziert.

In der Praxis führen diese Kosten dazu, dass die Optimizer-Schrittzeit eines verteilten Muon die kombinierte Wanduhrzeit der Forward- und Backward-Passes erreicht oder sogar übersteigt, was die algorithmischen Vorteile von Muon zunichtemacht. Dieser Malus ist besonders gravierend beim Training von Embodied Models (z. B. VLA), wo die Forward-Backward-Passes im Vergleich zum Standard-LLM-Pretraining rechentechnisch leichter sind als der Optimizer-Schritt.

2. Methodik: DMuon Systemdesign

DMuon ist eine Open-Source-Implementierung für verteiltes Training, die darauf ausgelegt ist, die Performance-Lücke zwischen Muon und AdamW zu schließen. Dies wird durch eine ko-designte Systemarchitektur erreicht, die Kommunikationsoptimierung, Lastverteilung und Kernel-Effizienz integriert, ohne das zugrunde liegende Trainingsframework (z. B. PyTorch/FSDP2) zu modifizieren.

2.1 Owner-zentrierte Ausführungsstrategie

Anstatt die Berechnung über alle Ranks zu replizieren, weist DMuon jedem Matrix-Parameter einen einzelnen Owner-Rank zu.

Rolle des Owners: Der Owner verwaltet den autoritativen Optimizer-Zustand und berechnet das Muon-Update (Newton–Schulz-Iteration) genau einmal pro Matrix.
Rolle der Non-Owner: Andere Ranks materialisieren den Parameter nur, wenn er für die Forward/Backward-Berechnung erforderlich ist, und führen die Orthogonalisierung nicht durch.
Kommunikationsfluss:
- Forward: Owners veröffentlichen aktualisierte Parameter asynchron an Non-Owners, gepipelint mit der Layer-Ausführung.
- Backward: Gradienten werden in gepackten Puffern akkumuliert und zum Owner reduziert.
- Overlap: DMuon nutzt eine zweistufige Hierarchie (Intra-Node und Inter-Node) mit feingranularen Gewichtslayouts, um die Kommunikation mit der Berechnung zu überlappen und kollektive Kontention zu verhindern.

2.2 Effiziente Gram Newton–Schulz Iteration

Um die Rechenkosten auf der Owner-Seite zu senken, verwendet DMuon die Gram Newton–Schulz Formulierung:

Gram-Raum-Rekursion: Anstatt im ursprünglichen $m \times n$ Raum zu operieren, wird die Iteration im $m \times m$ Gram-Raum ( $G = XX^\top$ ) durchgeführt. Dies reduziert die Komplexität von $O(m^2n)$ auf $O(m^3)$ , wenn $m < n$ .
Ausnutzung der Symmetrie: Da $G$ symmetrisch ist, nutzt DMuon einen SYRK-Stil Ausführungspfad, um nur den unteren Dreiecksanteil zu berechnen, was die arithmetische Arbeit der dominanten Matrixmultiplikation halbiert.
Batching: Um Probleme der GPU-Occupancy bei kleinen Matrizen zu adressieren, bündelt (batched) DMuon Matrizen ähnlicher Formen und führt die Gram-NS-Rekursion gemeinsam aus, um den Kernel-Dispatch-Overhead zu amortisieren.

2.3 Berechnungsorientierte Lastverteilung

Da die Kosten für Muon-Updates signifikant mit der Matrixform variieren, führt eine naive Zuweisung (z. B. Round-Robin) zu Stragglern (Nachzüglern). DMuon formuliert die Owner-Zuweisung als Mixed-Integer Linear Program (MILP):

Profiling: Während der Initialisierung benchmarkt das System die Ausführungszeit verschiedener Batch-Größen und Kernel-Schedules für jede distinkte Matrixform.
Optimierung: Der MILP-Solver minimiert die vorhergesagte Makespan (maximale Last) über die Owner-Ranks hinweg, um sicherzustellen, dass die schwersten Matrizen zur Lastverteilung verteilt werden.
Fallback: Für extrem große Suchräume wird eine Greedy-Heuristik verwendet, um die Initialisierungszeit zu begrenzen.

2.4 Form-adaptive Kernel-Autotuning

DMuon nutzt einen DSL-getriebenen Autotuning-Workflow (unter Verwendung von TileLang und CUTE DSL), um Hochleistungs-Kernel zu generieren.

Es sucht über Tile-Shapes, Block-Größen und Software-Pipeline-Konfigurationen, die spezifisch auf die Zielhardware und Matrixformen zugeschnitten sind.
Die optimale Konfiguration wird gecached und wiederverwendet, um sicherzustellen, dass der Kernel-Schedule immer nahezu optimal für die spezifischen Problemdimensionen im Transformer-Training ist.

3. Kernbeiträge

Feingranulare Kommunikationsoptimierung: Eine Owner-zentrierte Laufzeitumgebung, die den Kommunikations-Overhead durch hierarchisches Scheduling und feingranulare Gewichtslayouts minimiert und dabei die exakten Optimizer-Semantiken bewahrt.
Form-adaptiver Ausführungsstack: Eine Hochleistungs-Engine, die gebatchte Gram Newton–Schulz, Symmetrie-bewusste Kernel und Autotuning kombiniert, um heterogene Matrix-Workloads effizient zu handhaben.
Berechnungsorientierte Lastverteilung: Ein Messkosten-Optimierungsansatz (MILP), der Matrix-Owner basierend auf tatsächlichem Runtime-Profiling statt theoretischer FLOP-Zahlen zuweist.
Drop-in Integration: Ein Modul, das sich mit minimalem Aufwand in bestehende Trainings-Pipelines (z. B. FSDP2) integrieren lässt, ohne Änderungen am Framework selbst zu erfordern.
Validierung im Produktionsmaßstab: Erfolgreiche Implementierung sowohl bei Robotik- (Embodied) als auch bei Sprachmodell-Workloads.

4. Experimentelle Ergebnisse

Die Autoren evaluierten DMuon auf vier Produktions-Workloads: Wall-OSS-0.5, Pi0, Wall-WM und Qwen2.5-7B, skaliert von 8 auf 256 A800 GPUs.

Optimizer Step Time: DMuon erreicht eine 6,85- bis 163,00-fache Beschleunigung der Optimizer-Schrittzeit im Vergleich zu einer naiven „Gather-then-Compute“-Implementierung von Muon.
End-to-End Step Time: DMuon erzielt eine 1,48- bis 3,01-fache Beschleunigung der gesamten Schrittzeit.
Overhead vs. AdamW: Die Latenz pro Schritt von DMuon wird auf nahe AdamW-Niveaus gebracht, mit einem durchschnittlichen Overhead von nur +2 % gegenüber AdamW über die evaluierten Workloads hinweg.
Skalierung: Auf Wall-OSS hält DMuon den Durchsatz nahe an AdamW, während das Durchsatzwachstum bei steigender GPU-Anzahl bei naivem Muon durch schwere Kommunikations-Engpässe leiden würde.

Aufschlüsselung der Beschleunigung:

Symmetrischer Gram-Kernel: 48 % der Beschleunigung (Halbierung der Arithmetik in dominanten Produkten).
Owner-Scheduling & Lastverteilung: 32 % (Entfernung der $D$ -fachen redundanten Berechnung).
Autotuning & Batching: 16 % (Reduzierung des Kernel-Launch-Overheads für kleine Matrizen).

5. Bedeutung und Ansprüche

Das Paper behauptet, dass DMuon den Systems-Mismatch zwischen matrix-bewussten Optimierern und verteilten Trainingsinfrastrukturen erfolgreich löst. Durch die Eliminierung redundanter Berechnungen und die Optimierung der Kommunikation ermöglicht DMuon, dass die Konvergenzvorteile von Muon (wie höhere Recheneffizienz und bessere Optimierungsgeometrie) in praktische Wanduhr-Zeit-Beschleunigungen überführt werden.

Die Autoren positionieren DMuon nicht als neuen Algorithmus, sondern als notwendigen Systems-Enabler, der Muon zu einem praktikablen Drop-in-Ersatz für AdamW im groß angelegten Produktionstraining macht. Die Arbeit zeigt, dass der „Malus“ der matrix-bewussten Optimierung auf ein vernachlässigbares Niveau reduziert werden kann, was ein effizientes Scaling sowohl für Large Language Models als auch für Embodied Foundation Models (VLA) ermöglicht.

Das Paper stellt explizit klar, dass DMuon die Update-Regel oder die Konvergenzeigenschaften von Muon nicht ändert; es adressiert strikt den Systems-Overhead. Folglich sind die Vorteile in verteilten Settings am ausgeprägtesten, während das Single-GPU-Training eine moderatere Verbesserung sieht (ca. 2× Beschleunigung des Optimizers) aufgrund des Fehlens der zu eliminierenden verteilten Redundanz.