MXNorm: Reusing MXFP block scales for efficient tensor normalisation

Die Arbeit stellt MXNorm vor, eine effiziente Alternative zu RMSNorm, die durch Wiederverwendung von MXFP8-Blockskalen die Reduktionsgröße für die Normalisierung um den Faktor 32 verringert und dabei in Llama-3-Modellen nur minimale Genauigkeitsverluste bei signifikanten Geschwindigkeitssteigerungen erzielt.

Das Wesentliche

🚀 MXNorm: Der clevere Trick für schnellere KI-Modelle

Stell dir vor, du betreibst eine riesige Logistikfirma (das ist dein KI-Modell), die Millionen von Paketen (Daten) jeden Tag sortiert und weiterleitet.

1. Das Problem: Der Flaschenhals

In den letzten Jahren haben die LKWs für den Haupttransport (die Matrix-Multiplikation, also das eigentliche Rechnen) unglaublich schnell werden. Sie sind jetzt 80-mal schneller als früher! 🚛💨

Aber es gibt ein Problem: Die Verladestationen (die Normalisierungsschritte wie RMSNorm), an denen die Pakete vor dem Weitertransport gewogen, sortiert und auf die richtige Größe gebracht werden, sind nicht schneller geworden.

• Die LKWs warten oft auf die Verladestation.
• Die Verladestationen müssen jedes Paket einzeln wiegen, was viel Zeit und Energie kostet.
• Das bremst die ganze Firma aus, obwohl die LKWs bereit sind.

2. Die alte Lösung: RMSNorm (Der mühsame Wäger)

Bisher wurde jedes Paket einzeln gewogen, bevor es in den LKW geladen wurde.

• RMSNorm berechnet für jeden Daten-Satz die genaue „Durchschnittsgröße" (den Root Mean Square).
• Das ist sehr genau, aber es ist wie ein extra Schritt: Erst wiegen, dann sortieren, dann laden.
• Wenn du jetzt versuchst, die Pakete in winzige, leichte Kartons zu verpacken (um Speicherplatz zu sparen, genannt MXFP8), musst du sie trotzdem erst einzeln wiegen, bevor du sie in die neuen Kartons packst. Das ist doppelt gemoppelt.

3. Die neue Lösung: MXNorm (Der clevere Schachzug)

Die Forscher von Graphcore haben einen genialen Trick erfunden: MXNorm.

Stell dir vor, beim Verpacken in die neuen, leichten Kartons (MXFP8) muss man ohnehin ein Etikett anbringen, das sagt: „Dieser Karton wiegt maximal X Kilogramm" (das nennt man den Block Scale).

MXNorm sagt: „Warum sollen wir das Paket noch einmal extra auf der Waage wiegen? Wir nutzen einfach das Gewicht, das wir schon gerade erst berechnet haben, um das Etikett zu kleben!"

• Der Trick: Anstatt eine separate Waage (RMSNorm) zu benutzen, nutzen sie die Information aus dem Etikett (den Block-Scale), das ohnehin für die Komprimierung nötig ist.
• Das Ergebnis: Sie sparen sich den ganzen extra Wiege-Prozess. Die Verladestation wird um 32-mal kleiner und viel schneller.

4. Warum funktioniert das? (Die Analogie)

Stell dir vor, du hast eine Kiste voller Äpfel.

• RMSNorm: Du nimmst jeden Apfel einzeln, wiegst ihn, berechnest das Durchschnittsgewicht und normalisierst dann alle.
• MXNorm: Du weißt, dass die Kiste in Blöcke zu je 32 Äpfeln unterteilt ist. Für jeden Block hast du bereits das schwerste Äpfelchen notiert, um die Kiste sicher zu verpacken.
• Die Forscher haben mathematisch bewiesen: Wenn du das Gewicht des schwersten Äpfelchens in einem Block kennst, kannst du das Durchschnittsgewicht der gesamten Kiste ziemlich genau schätzen.
• Sie nutzen also das „Schwerste-Äpfelchen" (den Block-Scale), um die Waage zu ersetzen. Das ist fast genauso genau, aber viel schneller.

5. Die Ergebnisse: Schnell und stabil

Die Forscher haben diesen Trick an riesigen KI-Modellen (wie Llama 3) getestet:

• Geschwindigkeit: Die Modelle laufen auf neuer Hardware bis zu 2,4-mal schneller bei diesen spezifischen Schritten. Das klingt nach wenig, aber bei Milliarden von Berechnungen summiert sich das zu einem riesigen Zeitgewinn.
• Genauigkeit: Die KI lernt genauso gut wie vorher. Sie macht keine Fehler, nur weil wir die Waage durch einen Trick ersetzt haben.
• Stabilität: Sie haben herausgefunden, dass man nicht jeden Trick anwenden darf. Wenn man zu grob schätzt (z. B. nur den Durchschnitt der schwersten Äpfel nimmt), wird die KI instabil und macht Fehler. Aber mit der richtigen mathematischen Formel (dem „p=2"-Trick) klappt es perfekt.

🎯 Zusammenfassung für den Alltag

Stell dir vor, du musst 1000 Briefe in Umschläge stecken.

• Alt: Du wiegst jeden Brief einzeln, schreibst das Gewicht auf, und steckst ihn dann in den Umschlag.
• Neu (MXNorm): Du weißt, dass du für den Umschlag ohnehin das maximale Gewicht des schwersten Briefes im Stapel notieren musst. Du nutzt diese Notiz, um zu entscheiden, wie du die anderen Briefe behandelst. Du sparst dir das Wiegen jedes einzelnen Briefes.

Das Fazit: MXNorm ist ein „Drop-in"-Ersatz. Man muss das KI-Modell nicht neu erfinden, sondern tauscht einfach einen langsamen Baustein gegen einen schnelleren aus, der denselben Job erledigt, aber weniger Energie und Zeit verbraucht. Das ist ein großer Schritt, um KI-Modelle in Zukunft noch effizienter und günstiger zu machen.

Technische Zusammenfassung

1. Problemstellung

Die Leistungsfähigkeit von Deep-Learning-Workloads wird traditionell durch Matrixmultiplikationen limitiert. In den letzten acht Jahren haben Beschleuniger (GPUs) die Leistung bei niedrigen Präzisions-Matrixmultiplikationen (z. B. FP8, FP4) drastisch verbessert (Faktor 80x).

Allerdings haben sich andere Operationen, die für das Training essenziell sind – insbesondere Reduktionen (wie das Berechnen von Mittelwerten oder RMS) und elementweise Operationen – nur langsam verbessert (Faktor ~5x bis 9x). Da diese Operationen oft durch die Speicherbandbreite limitiert sind und nicht durch die Rechenleistung der Tensor-Cores, entstehen sie zu neuen Engpässen, sobald die Matrixmultiplikation optimiert ist.

Insbesondere die Normalisierungsschichten (wie RMSNorm in LLMs wie Llama) erfordern eine vollständige Reduktion über die versteckte Dimension (Hidden Dimension), um den Root Mean Square (RMS) jedes Tokens zu berechnen. Dies erfordert hohe Speicherbandbreite und verhindert eine vollständige Fusion mit nachfolgenden Quantisierungsschritten (MXFP), was die Effizienz mindert.

2. Methodik: MXNorm

Die Autoren schlagen MXNorm vor, einen Drop-in-Ersatz für RMSNorm, der die Normalisierung mit der Konvertierung in Microscaling-Formate (MXFP, z. B. MXFP8) fusioniert.

Kernidee:
Anstatt den RMS-Wert separat zu berechnen, nutzt MXNorm die bereits während der MXFP-Quantisierung berechneten Block-Skalen (Block Scales), um den RMS-Wert zu schätzen.

• Hintergrund (MXFP): Bei der MXFP-Quantisierung wird ein Tensor in Blöcke unterteilt. Für jeden Block wird ein Skalierungsfaktor (Scale) berechnet, der typischerweise auf dem Block-Absmax (dem absoluten Maximalwert im Block) basiert. Dieser Scale wird verwendet, um die Werte in ein niedrigpräzises Format (z. B. E4M3) zu skalieren.
• Approximation: Die Autoren beweisen mathematisch (Theorem 1), dass der generalized $p$-Mean der Block-Absmax-Werte linear mit dem RMS des gesamten Tensors skaliert, sofern die Daten einer Skalenfamilie (z. B. Gauß-Verteilung) folgen.
  • Formel für die Schätzung des inversen RMS ($\tilde{\rho}$):
    $$\tilde{\rho}_t = \tilde{c}(p, B) \cdot \left( \frac{1}{K} \sum_{k=1}^{K} \tilde{m}_{tk}^p \right)^{-1/p}$$
    Dabei ist $\tilde{m}_{tk}$ das Absmax des $k$-ten Blocks, $K$ die Anzahl der Blöcke, $B$ die Blockgröße und $\tilde{c}(p, B)$ ein Korrekturfaktor, der mittels Monte-Carlo-Simulationen für Gauß-Verteilungen bestimmt wurde.
• Fusion: Durch diese Schätzung können Normalisierung und Quantisierung in einem einzigen Durchgang über den Tensor erfolgen. Es wird nur eine einzige Statistik-Sammlung (über die Blöcke) benötigt, anstatt eine separate Reduktion für den RMS durchzuführen.

Unterschiede zu RMSNorm:

• RMSNorm: Berechnet den exakten RMS über alle $D$ Elemente ($\sqrt{\frac{1}{D}\sum x^2}$).
• MXNorm: Schätzt den RMS basierend auf den Block-Absmax-Werten.
• Stabilität: Die Autoren zeigen, dass die Wahl des Exponenten $p$ kritisch ist. $p=2$ (Quadratmittelwert der Absmaxes) liefert eine stabilere Schätzung und engere Obergrenzen für die Ausgabe als $p=1$ (arithmetisches Mittel), was für das Training großer Modelle entscheidend ist.

3. Wichtige Beiträge

1. Algorithmische Innovation: Einführung von MXNorm, das die Block-Skalen der MXFP-Quantisierung wiederverwendet, um die Normalisierung zu fusionieren. Dies reduziert die Größe der benötigten Reduktionsoperationen um den Faktor 32 (abhängig von der Blockgröße).
2. Theoretische Fundierung: Beweis, dass der generalized $p$-Mean der Block-Absmax-Werte asymptotisch gegen den RMS konvergiert (bis auf einen konstanten Faktor).
3. Implementierung: Entwicklung einer effizienten Kernel-Implementierung (MXNormLinear), die die Normalisierungsgewinne ($\gamma$) direkt in die nachfolgende Gewichtsmatrix integriert, um Speicher-Overhead zu vermeiden.
4. Analyse der Stabilität: Detaillierte Untersuchung von „Loss Spikes" (Trainingsausreißern). Es wird gezeigt, dass MXNorm mit $p=1$ bei großen Modellen (8B Parameter) instabil wird, da die Obergrenze der normalisierten Ausgabe zu hoch ist ($O(K)$ statt $O(\sqrt{K})$), was zu explodierenden Werten führt. MXNorm mit $p=2$ löst dieses Problem.

4. Ergebnisse

Die Methode wurde am Pre-Training von Llama 3-Modellen (125M, 1B und 8B Parameter) auf dem SlimPajama-Datensatz validiert.

• Trainingsstabilität und Genauigkeit:
  • Bei kleinen Modellen (125M, 1B) zeigen sowohl $p=1$ als auch $p=2$ kaum Unterschiede zu RMSNorm.
  • Bei großen Modellen (8B) führt $p=1$ zu Instabilitäten und schlechterem Final Loss.
  • MXNorm ($p=2$) erreicht jedoch einen Trainingsverlust und eine Zero-Shot-Leistung (OLMES-Benchmarks), die mit dem Baseline-RMSNorm (MXFP8) praktisch identisch ist.
• Performance-Gewinne:
  • Durch die Fusion von Normalisierung und Quantisierung und die Reduktion der Reduktionsoperationen wurden Kernel-Speedups von bis zu 2,4x für isolierte MXNorm-Kerne im Vergleich zu getrenntem RMSNorm + MXCast erzielt.
  • Auf Systemebene (Llama 3 8B Transformer-Layer) resultiert dies in einem Speedup von 1,3% für MXFP8 und 2,6% für NVFP4.
  • Dies unterstreicht die These, dass bei niedrigeren Präzisionsformaten die Optimierung von Nicht-Matrixmultiplikations-Operationen immer wichtiger wird.

5. Bedeutung und Fazit

MXNorm adressiert einen kritischen Engpass in der Skalierung von Large Language Models (LLMs) auf zukünftige Hardware-Architekturen.

• Effizienz: Es eliminiert redundante Speicherzugriffe und Reduktionsoperationen, die bei niedrigen Präzisionsformaten (FP8, FP4) den Flaschenhals darstellen.
• Praktikabilität: Die Methode ist ein „Drop-in"-Ersatz, der keine zusätzlichen Hyperparameter erfordert und mit bestehenden Frameworks (TorchAO, TorchTitan) kompatibel ist.
• Zukunftssicherheit: Da die Lücke zwischen Matrixmultiplikationsleistung und Reduktionsleistung weiter wächst, bietet MXNorm einen Weg, um die Effizienz von Normalisierungsschichten an die Leistungsfähigkeit moderner Tensor-Cores anzupassen.

Die Arbeit zeigt, dass durch intelligente Wiederverwendung von Quantisierungsstatistiken (Block Scales) signifikante Geschwindigkeitsgewinne erzielt werden können, ohne die Trainingsqualität oder Modellstabilität zu beeinträchtigen, sofern die richtigen mathematischen Approximationen ($p=2$) gewählt werden.