Practical FP4 Training for Large-Scale MoE Models on Hopper GPUs

Deze paper introduceert een trainingsmethode die MXFP4-activaties en expert-parallelle communicatie mogelijk maakt voor grote MoE-modellen op Hopper-GPU's zonder native FP4-ondersteuning, waardoor het piekgeheugengebruik met 14,8% en de doorvoer met 12,5% verbetert ten opzichte van FP8-baselines.

De kern

Stel je voor dat je een gigantische bibliotheek bouwt, niet met boeken, maar met kennis. Deze bibliotheek is zo groot dat hij 671 miljard "woorden" (parameters) bevat. Dit is een Groot Taalmodel (LLM), en het is een van de slimste kunstmatige intelligenties die we hebben.

Het probleem? Om deze bibliotheek te bouwen, heb je een enorm magazijn nodig (het geheugen van je computer) en een super-snelle vrachtwagen om de boeken heen en weer te rijden (de communicatie tussen de chips). Helaas is dit magazijn vol en de vrachtwagens steken vast in de file.

Deze paper vertelt het verhaal van een slimme oplossing die dit probleem oplost, zelfs zonder de allernieuwste, duurste vrachtwagens.

Het Probleem: De "Te Zware" Lading

Om zo'n enorme bibliotheek te trainen, gebruiken computers meestal een zware, nauwkeurige meeteenheid (noem het BF16 of FP8). Het is alsof je elke pagina van je boek in een zware, beschermde koffer verpakt voordat je hem op de vrachtwagen laadt.

• Nadeel: De vrachtwagens (de GPU's) worden snel overbelast. Ze kunnen niet genoeg boeken tegelijk vervoeren, en het magazijn raakt vol.
• De droom: Iedereen wil de "FP4"-stand gebruiken. Dit is alsof je de boeken in ultra-lichte, dunne plastic zakjes stopt. Je kunt er twee keer zoveel in een vrachtwagen proppen en het magazijn wordt half zo groot.
• Het obstakel: De huidige vrachtwagens (de Hopper-chips van NVIDIA) zijn niet gebouwd om die lichte zakjes direct te verwerken. Ze kunnen alleen de zware koffers aan. Als je de lichte zakjes probeert te gebruiken, moet je ze eerst weer in zware koffers verpakken, vervoeren, en weer uitpakken. Dat kost meer tijd dan het besparen!

De Oplossing: De Slimme "Tussenpersoon"

De auteurs van dit paper hebben een slimme truc bedacht. Ze hebben een software-gebaseerde vertaler gemaakt die werkt op de bestaande vrachtwagens, zonder dat je nieuwe hardware hoeft te kopen.

Hier is hoe het werkt, in drie simpele stappen:

1. De "Verpakking" (Quantisatie)

In plaats van de zware koffers te gebruiken, verpakken ze de boeken direct in die lichte FP4-zakjes.

• De truc: Ze doen dit heel slim. Ze gebruiken een speciale "verpakkingsmachine" (een computerprogramma) die de boeken direct in de lichte zakjes stopt, zonder eerst naar een zware koffer te gaan.
• Het resultaat: De vrachtwagens (de communicatie tussen de chips) zijn nu half zo zwaar. Ze kunnen veel sneller rijden.

2. De "Werkplek" vs. Het "Magazijn"

Dit is het meest creatieve deel.

• Op de werkplek (Berekening): Wanneer de vrachtwagen daadwerkelijk aan het werk is (rekenen), gebruiken ze nog steeds de zware koffers (FP8). Waarom? Omdat de rekenkracht van de machine daar het beste mee werkt.
• In het magazijn (Opslag): Zodra de boeken klaar zijn om opgeslagen te worden, verpakken ze ze direct in de lichte FP4-zakjes.
• De analogie: Stel je voor dat je een kok bent. Je kookt met zware, dure pannen (FP8) omdat dat het beste resultaat geeft. Maar zodra je het eten in de koelkast zet, doe je het in heel kleine, lichte plastic bakjes (FP4). Je hoeft niet met de zware pannen door de gang te lopen; je draagt alleen de lichte bakjes. Zo bespaar je ruimte en energie, zonder dat het eten (de kwaliteit) minder wordt.

3. De "Directe Vertaling" (Bit-voor-Bit)

Omdat de oude machines de lichte zakjes niet direct kunnen lezen, moesten ze een nieuwe vertaalcode schrijven.

• In plaats van: Lichte zakje -> Zware koffer -> Lichte zakje (wat traag is), hebben ze een code gemaakt die direct van Lichte zakje -> Zware koffer gaat.
• Ze hebben de "schalen" (de maten) van de lichte zakjes direct omgezet naar de maten van de zware koffers, zonder tussenstappen. Dit is alsof je een meetlat direct omzet van centimeters naar inches, zonder eerst naar meters te gaan.

Wat is het Resultaat?

Toen ze dit testten op hun gigantische bibliotheek (671 miljard parameters):

1. Ruimtebesparing: Ze hadden 15% minder ruimte nodig in het magazijn. Dit klinkt klein, maar bij zo'n groot systeem betekent dit dat ze veel meer boeken tegelijk kunnen verwerken.
2. Snelheid: Omdat de vrachtwagens minder zwaar hoefden te dragen, konden ze 12,5% sneller werken. In plaats van 1157 woorden per seconde, konden ze nu 1302 woorden per seconde verwerken.
3. Geen kwaliteitsverlies: De boeken werden niet beschadigd. De bibliotheek leerde net zo goed als met de oude, zware methode.

Conclusie

Deze paper laat zien dat je niet altijd de allernieuwste, duurste hardware nodig hebt om super-snel te zijn. Met een slimme software-uitvinding (een soort "tussenpersoon" die de lichte zakjes slim verpakt en vertaalt), kun je de bestaande machines veel efficiënter laten werken.

Het is alsof je een oude, zware vrachtwagen hebt, maar door slimme verpakkingstechnieken toch net zo snel en efficiënt bent als met een nieuwe, dure vrachtwagen. Dit maakt het trainen van de slimste AI's van de wereld veel goedkoper en sneller.

Technische samenvatting

Probleemstelling

Het trainen van grote-scale Mixture-of-Experts (MoE) modellen wordt momenteel beperkt door twee factoren: het geheugenvolume voor activaties (activation memory) en de communicatiekosten tussen experts (expert-parallel communication). Hoewel FP4 (4-bit floating point) formaten zoals MXFP4 en NVFP4 beloven om deze beperkingen te verhelpen, zijn ze op de huidige generatie NVIDIA Hopper GPU's (zoals de H100) niet direct bruikbaar.

De reden hiervoor is dat Hopper GPU's geen native ondersteuning bieden voor FP4-berekeningen of FP4-communicatieprimitieven. Bestaande pipelines zijn ontworpen voor BF16 of FP8. Het introduceren van FP4 in deze pipelines zonder native hardware-ondersteuning leidt tot inefficiënte "precision round-trips" (bijv. FP4 $\leftrightarrow$ BF16 $\leftrightarrow$ FP8), wat extra latentie, geheugenoverhead en precisieverlies veroorzaakt. Dit maakt FP4-training op Hopper tot nu toe onpraktisch, ondanks de grote potentie voor geheugen- en bandbreedtebesparing.

Methodologie

De auteurs presenteren een hybride precisie-trainingsframework dat MXFP4-efficiëntie mogelijk maakt op Hopper GPU's via software-emulatie, zonder native 4-bit Tensor Core-ondersteuning. De kern van de aanpak is het ontkoppelen van opslagprecisie (FP4) en rekenprecisie (FP8).

Belangrijkste technische componenten:

1. Hybride Precisie Dataflow:

   • Forward Pass: Activaties worden direct voor de All-to-All (A2A) communicatie gecomprimeerd naar MXFP4. Dit vermindert het communicatievolume en het geheugengebruik aanzienlijk. De kernberekeningen (GEMM) blijven echter in FP8 uitgevoerd om numerieke stabiliteit te waarborgen.
   • Backward Pass: Om de overhead van extra de-quantisatie te vermijden, wordt in de backward pass teruggegrepen naar de standaard FP8-communicatie. Dit asymmetrische ontwerp (agressieve FP4 in forward, conservatief FP8 in backward) optimaliseert de totale doorvoer zonder convergentie te schaden.

2. Directe Bitwise Conversie (FP4 $\to$ FP8):

   • In plaats van via BF16 te converteren (wat traag is), ontwikkelen de auteurs een directe bit-wise conversie-algoritme.
   • Dit algoritme pakt de sign, exponent en mantissa uit het FP4-bitstream en past deze direct toe op het FP8-formaat.
   • Een uitdaging is het verschil in blokgrootte: MXFP4 gebruikt blokken van 32 elementen, terwijl FP8 blokken van 128 gebruikt. De auteurs lossen dit op met een hiërarchische schaal-alignment strategie, waarbij de schaal van vier FP4-blokken wordt gemapt naar één FP8-blok.

3. Speciale CUDA Kernels:

   • Er zijn geoptimaliseerde kernels ontwikkeld voor layout-aware conversie:
     • BF16ToFP4Row: Quantisatie en packing tijdens de forward pass.
     • FP4RowToFP8Row: De-quantisatie voor standaard GEMM.
     • FP4RowToFP8Col: Een gefuseerde kernel die de-quantisatie combineert met matrix-transpositie voor gewichtsgradiënten (Wgrad), wat geheugentraffic minimaliseert.
   • De kernels ondersteunen ook "ragged tensors" (variabele lengte tokens per expert), wat essentieel is voor MoE-architecturen.

4. Communicatie-Optimalisatie:

   • De A2A-communicatie in de DeepEP-bibliotheek is uitgebreid om gepakte FP4-formaten te ondersteunen.
   • Door twee FP4-waarden in één byte te packen en schaal-factoren als 8-bit integer op te slaan, wordt de communicatiebelasting met ongeveer 50% verlaagd ten opzichte van FP8.

Belangrijkste Bijdragen

• FP4 Communicatie- en Cachingstrategie: Een strategie voor expert-parallelle MoE-lagen die het activatiegeheugen en het verkeer tussen GPU's met meer dan 50% reduceert.
• Directe Conversie-algoritme: Een lossless bit-wise conversie van FP4 naar FP8 met hiërarchische schaal-uitlijning, wat de noodzaak van BF16-intermediairs elimineert.
• Layout-bewuste Kernels: Geoptimaliseerde CUDA-kernels die quantisatie, dispatch en herberekening (recomputation) in FP4 ondersteunen, inclusief native support voor niet-gecontiguëerde MoE-tensors.
• Productie-deplooiing: De eerste implementatie van software-gemuleerde MXFP4 voor MoE-training op Hopper GPU's op productieschaal (671B parameters).

Resultaten

De methode is getest op een 671B-parameter MoE-model (DeepSeek-V3 configuratie) op een cluster van 256 NVIDIA Hopper GPU's.

• Geheugenreductie: De piekactivatiegeheugengebruik is met 14,8% (11,8 GB) verminderd ten opzichte van sterke FP8-baselines.
• Doorvoerverbetering: Door het lagere geheugengebruik kan de scope van recomputatie (herberekening van lagen tijdens de backward pass) worden verkleind. Dit leidt tot een toename van de trainingsthroughput van 1157 naar 1302 tokens per GPU per seconde, een verbetering van 12,5%.
• Convergentie: De trainingsverliezen volgen die van de BF16-baseline nauwkeurig, met een relatieve afwijking van slechts +0,61% ten opzichte van BF16, wat aantoont dat de kwantisatie geen stabiliteitsproblemen veroorzaakt.
• Vergelijking: Bij een 236B-model toont de methode consistent geheugenwinst (6,9% - 11% minder dan BF16/FP8) en maakt het trainen mogelijk onder condities waar BF16 en FP8 uitvallen door "Out-of-Memory" (OOM) fouten.

Betekenis

Dit werk bewijst dat FP4-efficiëntie praktisch realiseerbaar is op wijdverbreide Hopper-hardware, zelfs zonder native 4-bit Tensor Core-ondersteuning. Door zorgvuldig software-hardware co-design (specifieke kernels, directe conversie en hybride dataflow), kunnen organisaties nu profiteren van de geheugen- en bandbreedtevoordelen van FP4 voor grote MoE-modellen. Dit opent de deur voor het trainen van nog grotere modellen of het gebruik van grotere batchgroottes op bestaande infrastructuur, voordat de volgende generatie hardware (Blackwell) met native FP4-ondersteuning volledig beschikbaar is.