MXNorm: Reusing MXFP block scales for efficient tensor normalisation

Dit paper introduceert MXNorm, een efficiëntere variant van RMSNorm die bestaande MXFP8-blokschalen hergebruikt om normalisatie te versnellen en de benodigde reductiegrootte met 32x te verkleinen, wat leidt tot aanzienlijke snelheidswinst bij het trainen van modellen zoals Llama 3 zonder noemenswaardig verlies aan nauwkeurigheid.

De kern

Het Grote Probleem: De Snelweg en de Tolpoortjes

Stel je voor dat het trainen van een slimme computer (zoals een AI die tekst schrijft) een gigantische vrachtwagen is die over een snelweg rijdt.

• De vrachtwagen (Matrixvermenigvuldiging): Dit is het zware werk van het rekenen. De technologie hier is de afgelopen jaren razendsnel geworden. De vrachtwagen kan nu 80 keer sneller rijden dan voorheen.
• De tolpoortjes (Normalisatie): Tussendoor moet de vrachtwagen echter vaak stoppen bij een tolpoortje om te controleren of de lading goed verdeeld is. Dit heet "normalisatie".

Het probleem is dat de tolpoortjes niet sneller zijn geworden. Terwijl de vrachtwagen razendsnel is, staan de tolpoortjes nog steeds in de file. De vrachtwagen moet wachten, waardoor de hele snelweg vastloopt. De onderzoekers van Graphcore zeiden: "We moeten die tolpoortjes sneller maken, of ze zelfs overbodig maken."

De Oplossing: MXNorm (De Slimme Tolcontrole)

In de huidige wereld van AI gebruiken ze een methode genaamd RMSNorm. Dit is een tolcontrole die heel nauwkeurig is, maar veel tijd kost. Ze moeten elke vrachtwagen (elk stukje data) apart afwegen en controleren.

Daarnaast gebruiken ze een nieuwe manier om data op te slaan, genaamd MXFP. Dit is als het verpakken van de vrachtwagen in compacte, lichte dozen. Om deze dozen te maken, moeten ze al een keer controleren hoe zwaar de zwaarste lading is in elke doos. Dit noemen ze "block scales".

Het inzicht van MXNorm:
De onderzoekers dachten: "Wacht even! We hebben al een schatting van het gewicht nodig om de dozen te maken. Waarom doen we dan nog een extra, dure controle (RMSNorm) voordat we de dozen maken? Laten we die twee stappen samenvoegen!"

MXNorm is dus een slimme truc waarbij ze de gewichten die ze al hebben berekend voor het verpakken van de data, direct ook gebruiken om de "tolcontrole" (normalisatie) te doen. Ze hoeven niet twee keer te tellen, maar slechts één keer.

Hoe werkt het in de praktijk?

1. De Oude Weg (RMSNorm + MXCast):

   • Stap 1: Tel precies op wat de gemiddelde zwaarte is van alles (duur en traag).
   • Stap 2: Pas de lading aan op basis van dat gemiddelde.
   • Stap 3: Verpak het in de kleine dozen (MXFP).
   • Resultaat: Veel wachttijd.

2. De Nieuwe Weg (MXNorm):

   • Stap 1: Kijk naar de zwaarste lading in elke doos (dit moet je al doen om de doos te verpakken).
   • Stap 2: Gebruik die zwaarste lading om een slimme schatting te maken van het gemiddelde gewicht.
   • Stap 3: Pas de lading aan en verpak het tegelijkertijd.
   • Resultaat: Geen extra wachttijd. Het is alsof je de tolpoortjes hebt vervangen door een scanner die al in de vrachtwagen zit.

Wat hebben ze ontdekt?

De onderzoekers hebben dit getest op verschillende maten van AI-modellen (kleine en grote hersenen).

• Snelheid: Omdat ze minder werk hoeven te doen, gaat het 1,3% tot 2,6% sneller. Klinkt misschien klein, maar bij een AI die urenlang rekent, scheelt dat veel tijd en energie.
• Stabiliteit: Ze ontdekten dat je niet zomaar elke schatting kunt gebruiken. Als je te simpel rekent (zoals een gemiddelde nemen), kan de AI soms "ontsporen" (zoals een vrachtwagen die van de weg raakt). Maar als je een iets slimmere schatting gebruikt (de "kwadratische gemiddelde" methode), werkt het net zo goed als de oude, zware methode.
• Kwaliteit: De AI leert even goed als voorheen. De resultaten zijn bijna identiek, maar dan veel sneller.

Waarom is dit belangrijk?

We gaan naar een toekomst waar computers nog sneller worden en data nog kleiner wordt verpakt (zoals van 8 bits naar 4 bits). In die wereld wordt de "tolcontrole" (normalisatie) nog belangrijker als een knelpunt.

MXNorm is als het vinden van een nieuwe route die de file omzeilt. Het maakt AI-trainingen goedkoper, sneller en efficiënter, zonder dat de slimheid van de AI eronder lijdt. Het is een slimme manier om bestaande gereedschappen (de schaal van de dozen) te hergebruiken voor een nieuw doel, zodat we niet meer twee keer hoeven te werken.

Kort samengevat:
MXNorm is een slimme truc die twee taken in één keer doet: het verpakken van data in kleine dozen én het controleren of die data goed is. Hierdoor staat de AI niet meer vast in de file, maar rijdt hij soepel door.

Technische samenvatting

1. Het Probleem

De prestaties van matrixvermenigvuldiging (matmul) in diepe leerwerkloads zijn de afgelopen jaren exponentieel verbeterd door nieuwe hardware-accelerators die lage precisie (zoals FP8 en FP4) ondersteunen. Echter, andere operationele onderdelen van neurale netwerken, zoals reducties (bijv. het berekenen van gemiddelden of RMS) en elementsgewijze berekeningen, hebben niet in hetzelfde tempo meegegroeid.

• Bottleneck: In moderne architecturen (zoals Pre-Norm Transformers, bijv. Llama 3) wordt normalisatie (RMSNorm) uitgevoerd voordat de data wordt omgezet naar lage precisie (MXFP).
• Inefficiëntie: RMSNorm vereist een volledige reductie over de verborgen dimensie om de wortel-gemiddelde-kwadraat (RMS) te berekenen. Dit gebeurt vaak nog in hogere precisie en vormt nu een beperkende factor voor de doorvoer, vooral omdat de bandbreedte voor geheugen en CUDA-kernen niet even snel is toegenomen als de Tensor Core-prestaties.
• Doel: Er is behoefte aan een normalisatiemethode die minder overhead veroorzaakt en beter integreert met de bestaande MXFP (Microscaling Floating Point) quantisatie-workflows.

2. Methodologie: MXNorm

De auteurs stellen MXNorm voor, een "drop-in" vervanging voor RMSNorm die de normalisatie fuseert met de conversie naar MXFP-formaten.

Kernidee:
In plaats van de RMS van de volledige tensor opnieuw te berekenen, benut MXNorm de block scales die al worden berekend tijdens het quantiseren van de tensor naar MXFP8 (Microscaling).

• Herkomst van de schaal: Bij MXFP quantisatie wordt de tensor opgesplitst in blokken. Voor elk blok wordt een absmax (het grootste absolute waarde) berekend en omgezet naar een schaalfactor (block scale) in een speciaal formaat (E8M0).
• De Approximatie: De auteurs bewijzen wiskundig (Theorema 1) dat de RMS van een tensor benaderd kan worden door een veralgemeende $p$-mean van de block absmaxes te nemen.
  • Als een vector lineair wordt geschaald, schalen zowel de RMS als de block absmaxes mee.
  • De verhouding tussen de RMS en de gemiddelde block absmax is constant voor een gegeven blokgrootte en verdeling (bijv. Gaussisch).
• Implementatie:
  • De RMS ($\rho$) wordt geschat als: $\tilde{\rho} = c \cdot (\frac{1}{K} \sum m_k^p)^{-1/p}$, waarbij $m_k$ de absmax van blok $k$ is en $c$ een correctiefactor is.
  • De auteurs testen twee varianten: $p=1$ (arithmetic mean) en $p=2$ (RMS van de absmaxes).
  • De normalisatie en de quantisatie worden samengevoegd in één pass over de tensor, waardoor de reductie-grootte met een factor 32x wordt verkleind (afhankelijk van de blokgrootte).

Gradienten en Training:
Om backpropagation mogelijk te maken, gebruiken de auteurs een "straight-through estimator". Ze hergebruiken de gradientenberekening van RMSNorm, maar cacheën de invoer en de geschatte inverse RMS om de gradienten in hoge precisie te berekenen zonder extra geheugenoverhead. De learnable gain parameter ($\gamma$) wordt gefuseerd met de gewichten van de daaropvolgende lineaire laag.

3. Belangrijkste Bijdragen

1. MXNorm Architectuur: Een nieuwe normalisatielaag die de reductie-overhead elimineert door bestaande quantisatie-schaalfactoren te hergebruiken.
2. Wiskundige Validatie: Een theoretisch bewijs dat de RMS benaderd kan worden via block absmaxes met een constante correctiefactor, geldig voor i.i.d. steekproeven uit een schaal-familie verdeling.
3. Stabiliteitsanalyse: Een diepgaande analyse van waarom bepaalde varianten (zoals $p=1$) falen bij grote schaal en waarom andere ($p=2$) stabiel blijven. Het paper toont aan dat de bovengrens van de genormaliseerde output cruciaal is voor stabiliteit; MXNorm met $p=2$ biedt een vergelijkbare bovengrens als RMSNorm, terwijl $p=1$ te losse grenzen heeft die leiden tot instabiliteit.
4. Praktische Implementatie: Een volledige PyTorch-implementatie die compatibel is met bestaande frameworks (TorchAO, TorchTitan) en werkt op commerciële hardware (NVIDIA GB200).

4. Resultaten

De methode is gevalideerd op het pre-trainen van Llama 3 modellen met 125M, 1B en 8B parameters.

• Training Stabielheid:
  • Bij kleine modellen (125M, 1B) presteert MXNorm ($p=1$ en $p=2$) vergelijkbaar met de RMSNorm-baseline.
  • Bij het 8B model faalt de $p=1$ variant: deze vertoont verlies-spike's (loss spikes) en instabiliteit. De oorzaak is dat $p=1$ een te hoge bovengrens toestaat voor extreme waarden (outliers), wat leidt tot explosieve updates.
  • De $p=2$ variant (gebruikmakend van de kwadratische mean van de absmaxes) bereikt een identiek verlies als de RMSNorm-baseline (2.126 vs 2.132) en toont geen instabiliteit.
• Zero-Shot Performance: Op OLMES-taken presteert MXNorm ($p=2$) vergelijkbaar met RMSNorm, met vergelijkbare scores op benchmarks zoals MMLU, ARC, en HellaSwag.
• Prestatie (Snelheid):
  • Door het samenvoegen van normalisatie en quantisatie, wordt de kernel sneller uitgevoerd.
  • Op NVIDIA GB200 hardware werd een 2.4x snelheidswinst waargenomen voor de geïsoleerde MXNorm-kernel ten opzichte van RMSNorm + MXCast.
  • Voor de volledige transformer-laag van Llama 3 8B resulteert dit in een 1.3% snelheidswinst voor MXFP8 en 2.6% voor NVFP4. Dit is significant gezien de focus op het optimaliseren van niet-matmul operaties bij lage precisie.

5. Significantie en Conclusie

MXNorm lost een kritiek probleem op in de schaalbaarheid van Large Language Models (LLMs) op toekomstige hardware. Terwijl matrixvermenigvuldiging steeds sneller wordt, blijven normalisatielagen een bottleneck.

• Efficiëntie: Het elimineert de noodzaak voor een aparte, dure reductie-operatie voor normalisatie, wat de geheugendruk verlaagt en de doorvoer verhoogt.
• Toekomstbestendigheid: De methode is niet beperkt tot FP8; het kan worden toegepast op nog lagere precisie formaten (zoals INT2 of ternair) en andere quantisatiemethoden die block absmaxes berekenen.
• Praktische Impact: Het biedt een eenvoudige, "drop-in" oplossing voor bestaande LLM-architecturen (zoals Llama 3) die de prestaties behoudt terwijl het de inferentie- en trainingskosten verlaagt op moderne accelerators.

Kortom, MXNorm demonstreert dat door slimme hergebruik van quantisatiemetagegevens, fundamentele operationele bottlenecks in AI-accelerators kunnen worden opgelost zonder in te leveren op modelkwaliteit.