NVIDIA hat Muon sowie weitere neuere Optimizer in Megatron Core integriert und beschreibt dazu Techniken, mit denen sich diese Verfahren für das Training großer Sprachmodelle auf viele GPUs verteilen lassen. Muon wurde zuletzt bereits für Open-Source-Modelle wie Kimi K2 und GLM-5 verwendet.
Nach den im Beitrag gezeigten Messungen liegt der Trainingsdurchsatz von Muon auf dem NVIDIA-GB300-NVL72-System nahe bei AdamW. Tabelle 1 fasst dafür Ergebnisse für Kimi K2 und Qwen3 30B zusammen. NVIDIA schreibt, dass der Trainings-Performanceverlust mit Muon gegenüber AdamW sehr klein sei. Die Model FLOPs Utilization sei mit Muon höher, wenn die FLOPs der Matrixmultiplikationen in den Newton-Schulz-Iterationen mitgezählt werden.
Die Messungen wurden mit NVIDIA NeMo Megatron Bridge 26.02 durchgeführt, einer PyTorch-nativen Bibliothek innerhalb des NeMo Frameworks für Pretraining, SFT und LoRA populärer LLM- und VLM-Modelle. Für Kimi K2 nutzte NVIDIA 256 NVIDIA-GB300-GPUs mit PP4DP64EP64, für Qwen3 30B-A3B acht NVIDIA-GB300-GPUs mit DP8EP8.
Der Beitrag benennt zugleich die praktischen Hürden solcher höherwertigen Optimizer. Der Preconditioning-Schritt mit Newton-Schulz-Iteration oder Eigenzerlegung erhöht Rechenaufwand und Speicherverbrauch deutlich. Hinzu kommen mögliche numerische Instabilitäten bei Mixed-Precision-Training und Gradient Accumulation sowie zusätzliche Kommunikationskosten, wenn synchronisierte und orthogonalisierte Updates über viele GPUs verteilt werden.
Um diese Probleme zu adressieren, führt NVIDIA einen layer-wise distributed optimizer ein. Während elementweise verteilte Optimizer wie bei AdamW Zustände und Gradienten gleichmäßig über GPUs aufteilen, benötigen Optimizer wie Muon für die Berechnung ihrer Updates die Gradienten einer vollständigen Schicht. Deshalb werden bei der layer-weisen Verteilung vollständige Layer verschiedenen Data-Parallel-Rängen zugewiesen, sodass jede GPU komplette Schichten für die Berechnung des Preconditioners besitzt.
Eine Folge dieser Verteilung ist Kommunikation mit variabler Größe, weil sich ganze Schichten in ihrer Größe unterscheiden. Dafür wird all_gatherv verwendet, um unterschiedlich große Parameter-Updates zwischen GPUs zusammenzuführen. Diese Implementierung ist laut NVIDIA vollständig in Megatron Core integriert; verwiesen wird auf die Datei layer_wise_optimizer.py.
Zusätzliche Herausforderungen entstehen bei Tensor Parallelism, weil dabei einzelne Gewichtsmatrizen entlang bestimmter Dimensionen über mehrere GPUs aufgeteilt werden und kein Gerät den vollständigen Parametertensor hält. Für den zentralen Orthogonalisierungsschritt von Muon wird jedoch die vollständige Momentum-Puffermatrix benötigt. NVIDIA beschreibt dafür in TensorParallelMuon zwei Modi.
Im duplicated mode werden die Momentums zunächst über die TP-Domäne per all-gather zusammengeführt, sodass jede GPU alle Daten für die Newton-Schulz-Iteration besitzt. Anschließend führt jede GPU die vollständigen Iterationen selbst aus und aktualisiert ihre eigenen Gewichte mit dem entsprechenden orthogonalisierten Momentum. Dieser Modus ist auf Netzwerklatenz optimiert, führt aber dazu, dass jede GPU dieselben Newton-Schulz-Iterationen ausführt. Pro Update ist unabhängig von der Zahl der Newton-Schulz-Iterationen nur eine all-gather-Kommunikation zu Beginn nötig.
Im distributed mode wird die Berechnung der Newton-Schulz-Iterationen über die GPUs der TP-Domäne verteilt. Jede Newton-Schulz-Iteration umfasst drei Matrixmultiplikationen; nach der ersten Matrixmultiplikation jeder Iteration ist ein all-reduce erforderlich. Dieser Modus ist auf Rechenaufwand optimiert, benötigt aber häufigere Kommunikation.
Zusätzlich unterstützt NVIDIA einen blockwise mode. Dabei erfolgen Orthogonalisierung und Update nur mit dem Momentum und den Gewichten, die jeweils lokal auf einer GPU liegen. Dieser Modus ist kostengünstiger zu berechnen und benötigt keine zusätzliche Kommunikation, ist aber nicht äquivalent zur Orthogonalisierung der gesamten Momentum-Matrix. NVIDIA beschreibt ihn als konzeptionell blockweise Orthogonalisierung, ähnlich dem Blocking-Ansatz aus Scalable Second Order Optimization for Deep Learning.
Als weitere Optimierungen nennt NVIDIA Communication Hiding, besseres Load Balancing sowie SYRK- und fused-all-reduce-Ansätze. Das Einsammeln von Parametern erfolgt in der aktuellen Version direkt nach dem Optimizer-Schritt und kann unter bestimmten Bedingungen zum Flaschenhals werden; künftig soll dieses Gathering bis zum Forward-Schritt des nächsten Batches verzögert und mit Berechnung überlappt werden. Beim Load Balancing verweist NVIDIA darauf, dass sich die Optimizer-Kosten je nach Größe von Attention- und MLP-Gewichtsmatrizen zwischen den Schichten unterscheiden und perfektes Lastgleichgewicht selbst theoretisch nicht erreichbar sei; derzeit werden Layer im Round-Robin-Verfahren verteilt.
Bei Newton-Schulz-Iterationen lassen sich laut NVIDIA die ersten zwei der drei Matrixmultiplikationen auf SYRK abbilden, wodurch nahezu die Hälfte der Floating-Point-Operationen eingespart werden könne; etwas weniger als die Hälfte müsse dennoch vollständig berechnet werden, weil Kacheln entlang der Diagonalen vollständig ausgerechnet werden müssen. Triton-Kernels sind dafür bereits enthalten. Für die Zukunft ist eine CuTe-DSL-Implementierung geplant, um Kommunikation in den SYRK-Kernel zu integrieren und Bandbreite sowie Kommunikationslatenz auf Tile-Ebene besser zu verbergen.
Neben Muon unterstützt NVIDIA nach eigenen Angaben weitere Optimizer für die Forschung, darunter MOP, das als orthogonalisierter Optimizer auf Polar Decomposition basiert, sowie eine SOAP-Variante, die die Eigenbasis pro Schritt mit Eigenzerlegung und KL-Korrektur in REKLS aktualisiert.
Zur Reproduzierbarkeit verweist NVIDIA auf die Performance Recipes im GitHub-Repository von NVIDIA-NeMo/Megatron-Bridge. Muon ist in Megatron Core als eigene Klasse sowie mit einem layer-wise distributed wrapper integriert. Für Muon leite use_distributed_optimizer automatisch auf den layer-weisen statt auf den elementweisen verteilten Optimizer um.
Für die Konfiguration nennt NVIDIA drei Punkte: Beim NS-Modus soll duplicated mode verwendet werden, wenn Netzwerklatenz dominiert, und distributed mode, wenn die Berechnung der Engpass ist; blockwise vermeidet Kommunikation vollständig, approximiert aber die vollständige Orthogonalisierung. Für große DP-Skalierung soll der layer-wise distributed optimizer Muon mit geringem Overhead ermöglichen. Außerdem nennt NVIDIA künftige Verbesserungen bei Communication Hiding, Load Balancing und fused SYRK/all-reduce, um verbleibende Durchsatzunterschiede weiter zu verringern.
