Technisch Overzicht: Full-Stack FP4: Stabiele LLM Pretraining met Gekwantiseerde Projecties, Optimizers en Attention

1. Probleemstelling

Recente vooruitgang in NVIDIA Blackwell GPU's heeft ondersteuning voor native NVFP4 (4-bit) Tensor Cores een praktisch doel gemaakt voor de pretraining van Large Language Models (LLM's). Terwijl bestaand onderzoek succesvol heeft aangetoond dat stabiele 4-bit pretraining mogelijk is voor Transformer lineaire projecties, blijven er significante hiaten bestaan in andere kritieke trainingsmodules. De meeste huidige NVFP4-recepten behouden een hogere precisie (bijv. BF16 of FP8) voor optimizer states, optimizer berekeningen en attention-mechanismen.

De primaire uitdaging is dat deze modules over verschillende numerieke structuren beschikken die een uniforme kwantiseringsregel niet kan adresseren:

• Lineaire Projecties: Gradiëntsignalen verzwakken tijdens de training terwijl kwantiseringsruis constant blijft, wat de effectiviteit van voortgezette lage-precisie training vermindert.
• AdamW States: Het tweede-moment momentum is niet-negatief, heavy-tailed, persistent en gevoelig voor denominator-waarden, waardoor directe 4-bit opslag gevoelig is voor distortie.
• Root/Muon Optimizers: Deze vertrouwen op herhaalde Newton–Schulz matrixproducten waarbij fouten in lage precisie kunnen accumuleren en convergentie kunnen destabiliseren.
• Attention: De operator bevat softmax-gevoelige paden (bijv. $PV$, $dOV^\top$) die hoge precisie vereisen, terwijl andere paden (bijv. $QK^\top$) robuuster zijn tegen kwantisatie.

Consequent is er een gebrek aan een verenigd, modulair framework dat stabiele, native NVFP4 pretraining over de gehele trainingsstack (projecties, optimizers en attention) mogelijk maakt zonder te vertrouwen op mixed-precision fallbacks voor kritieke componenten.

2. Methodologie: Full-Stack FP4

De auteurs presenteren Full-Stack FP4, een modulair framework bestaande uit vier onafhankelijke, composabele recepten die ontworpen zijn om de specifieke numerieke moeilijkheden van elke trainingsmodule aan te pakken.

2.1 LoRA-SVD Lineaire Projecties

Om de degradatie van gradiëntsignalen in latere trainingsfasen te beperken, stellen de auteurs LoRA-SVD voor.

• Mechanisme: Deze methode parametriseert gewichten als $W = W_{res} + \beta L_2 L_1$. Het behoudt een compacte hoofdsubruimte ($L_2 L_1$) in BF16 terwijl de dominante residu ($W_{res}$) wordt berekend in NVFP4.
• Backward Pass: Een subspace-geconstreineerde backward pass gebruikt een Cholesky-QR constraint om ervoor te zorgen dat de BF16 low-rank branch updates uitlijnen met de hoofdsubruimte, terwijl het residu de NVFP4 component bijwerkt.
• Realignment: De subruimte wordt periodiek opnieuw uitgelijnd (elke 2.048 stappen) met behulp van randomized SVD om de componenten te resetten, waardoor wordt gegarandeerd dat de BF16 subruimte optimaal blijft.

2.2 Quantized-State AdamW (Q-AdamW)

Om de opslag van persistente, heavy-tailed momentum states in NVFP4 te stabiliseren:

• Eerste Moment ($m_t$): Past Hadamard mixing toe om de distributie te verzachten voordat Stochastic Rounding (SR) en adaptieve 4/6 kwantisatie plaatsvindt.
• Tweede Moment ($v_t$): Maakt gebruik van een driefasige preprocessing pipeline voor NVFP4 opslag:
  1. Vierkantswortel (Square Root): Verkort de heavy tail.
  2. Tile Mean: Scheidt de positieve gemeenschappelijke component (opgeslagen in BF16) van een ongeveer nul-gemiddelde residu.
  3. Hadamard Transform: Verspreidt gelokaliseerde energie in het residu om het symmetrischer te maken.
• Reconstructie: De state wordt gereconstrueerd door de transformatie en de gemiddelde scheiding te inverteren, en vervolgens het resultaat te kwadrateren. Deze pipeline voorkomt de noodzaak van SVD-gebaseerde subspace operaties.

2.3 NVFP4 Root Optimizer (Q-Root)

Voor de Newton–Schulz (NS) iteraties die worden gebruikt in de Root optimizer:

• Directe NVFP4 Executie: In plaats van een mixed-precision route, worden de NS matrixproducten direct in NVFP4 uitgevoerd.
• Stabilisatie: Stabiliteit wordt bereikt door vormafhankelijke coëfficiënten (geoptimaliseerde $a, b, c$ voor specifieke matrix aspect ratio's) en p99 outlier clipping. Deze aanpak beperkt foutversterking over opeenvolgende NS-stappen zonder dat er hulpbronnen voor mixed-precision paden nodig zijn.

2.4 Mixed-Precision Attention (Q-Attn)

Om precisie en efficiëntie in attention te balanceren:

• Precisie Splitsing: Het framework wijst NVFP4 toe aan robuuste paden ($QK^\top$ en $dS$) en behoudt BF16 voor softmax-gevoelige paden ($PV$, $P^\top dO$, en $dOV^\top$).
• Consistentie: Gekwantiseerde views van $Q, K, V$ worden hergebruikt tussen de forward en backward passes om tensorconsistentie te waarborgen en redundant geheugentrafic te voorkomen.

3. Belangrijkste Bijdragen

1. LoRA-SVD Projecties: Een combineerbare herparametrisering die de lineaire-alleen verliesgap vermindert van 1,40% naar 0,61% vergeleken met standaard NVFP4, waardoor een compacte subruimte in BF16 wordt beschermd.
2. Quantized-State AdamW: De eerste demonstratie van stabiele NVFP4 AdamW momentum-state opslag in een geïntegreerde pretraining run, gebruikmakend van een getransformeerde state pipeline die de reconstructiefout van het tweede moment met ~60% vermindert.
3. NVFP4 Root: Een direct low-precision Newton–Schulz pad mogelijk gemaakt door vormbewuste coëfficiënten en clipping, wat de iteratiefout vermindert van ~10% naar ~3%.
4. Trainable Mixed-Precision Attention: Een strategie die softmax-gevoelige paden in BF16 beschermt terwijl geschikte forward en backward paden gekwantiseerd worden, wat een consistente NVFP4 pretraining operator mogelijk maakt.

4. Experimentele Resultaten

De auteurs evalueerden het framework op een 3B parameter Transformer getraind voor 64B tokens.

• Training Loss: De BF16 baseline bereikte een loss van 2,267, terwijl Full-Stack FP4 2,286 bereikte, wat een relatief verschil geeft van 0,838%.
• Downstream Performance: Zero-shot evaluatie toonde aan dat Full-Stack FP4 een gemiddelde perplexiteit had van 26,665 (vs. 26,675 voor BF16) en een gemiddelde nauwkeurigheid die slechts 0,10 procentpunt lager is dan de BF16 baseline.
• Ablatie Studies:
  • LoRA-SVD alleen verminderde de lineaire-alleen loss gap van 1,40% naar 0,61%.
  • De modulaire recepten kunnen onafhankelijk worden ingeschakeld of gecombineerd, waarbij "Full-Stack" de geaggregeerde dekking vertegenwoordigt.
• Native Hardware Efficiëntie (RTX 5090):
  • Root Speedup: Full-Stack FP4 bereikte 2,50×–2,83× versnellingen ten opzichte van geoptimaliseerde BF16 en 1,83×–2,17× ten opzichte van NVIDIA Transformer Engine (TE) voor de Root optimizer stap.
  • Geheugenbesparing: Het framework verminderde de AdamW piekgeheugengebruik met 37,9%–42,5% vergeleken met BF16.
  • Tijdswinst: Geschatte tijdswinst bij gelijkwaardige loss werd berekend op 56,9%–61,9% wanneer rekening wordt gehouden met de verminderde staptijd en de lichte toename in benodigde tokens om de BF16 loss te evenaren.

5. Betekenis en Beperkingen

Betekenis:
Het paper demonstreert dat een "Full-Stack" aanpak voor NVFP4 pretraining haalbaar is, waarbij projecties, optimizer states, optimizer berekeningen en attention worden gedekt. Het stelt vast dat native 4-bit pretraining een prestatie kan leveren die vergelijkbaar is met BF16 (binnen <1% loss gap), terwijl het aanzienlijke geheugen- en doorvoergvoordelen biedt op Blackwell-klasse hardware. Het werk levert een modulaire set recepten die het voor practitioners mogelijk maakt om laag-precisie stapsgewijs te adopteren.

Beperkingen:

• Schaal: Experimenten zijn beperkt tot modellen tot 3B parameters en 64B tokens. Het paper stelt geen gedrag vast bij schalen van 7B+ of over langere trainingshorizonten.
• Efficiëntie Meting: Native hardware efficiëntie werd gemeten op vier decoder blocks op een enkele RTX 5090, niet op een volledig gedistribueerd model.
• Schattingen: Tijdswinstberekeningen zijn cross-setting schattingen gebaseerd op vaste batchgroottes en houden geen rekening met potentiële doorvoertowinsten door het vergroten van de batchgrootte met vrijgekomen geheugen.
• Kernel Fusion: Sommige kernels (LoRA backward, attention) zijn nog niet volledig gefuseerd, en de implementatie maakt nog geen gebruik van alle FlashAttention-4 optimalisaties.
• Formaat Specificiteit: De recepten en metingen zijn specifiek voor NVFP4 en beweren geen overdraagbaarheid naar MXFP4.