Sparser, Faster, Lighter Transformer Language Models

Probleemstelling

De schaalvergroting van autoregressieve Large Language Models (LLMs) heeft tot ongeëvenaarde vooruitgang geleid, maar brengt enorme rekenkosten met zich mee. De feedforward-lagen (FFN) van een Transformer vormen de grootste bron van parameters en uitvoeringskosten (FLOPs). Hoewel onderzoek heeft aangetoond dat er van nature sprake is van activatiesparsity (veel neuronen zijn voor een bepaald token inactief) in deze lagen, is het benutten van deze ongestructureerde sparsity op moderne GPU's problematisch.

Er bestaat een fundamentele mismatch tussen ongestructureerde sparsity en GPU-architecturen, die geoptimaliseerd zijn voor dichte (dense) berekeningen. Traditionele sparse formaten (zoals ELLPACK) vereisen vaak extra synchronisatie, memory overhead voor indexbeheer en leiden tot heterogene werklasten. Hierdoor zijn officiële kernels voor sparse operaties vaak zelfs langzamer dan hun dichte tegenhangers, wat de praktische adoptie van sparsity voor efficiëntieverbetering heeft geblokkeerd.

Methodologie

De auteurs introduceren een nieuwe aanpak die bestaat uit drie pijlers: een nieuw dataformaat, aangepaste CUDA-kernels en een trainingsstrategie.

1. Nieuwe Sparse Formats:

• TwELL (Tile-wise ELLPACK): Voor inferentie wordt een nieuw formaat voorgesteld dat is ontworpen om naadloos te integreren met de gepipelinde uitvoering van moderne GPU's. In plaats van de hele rij te bekijken (zoals bij ELL), verdeelt TwELL de kolommen in horizontale tegels (tiles). Binnen elke tegel worden niet-nul waarden en hun indices lokaal opgeslagen. Dit maakt het mogelijk om het formaat direct te genereren in de epilogue van een geoptimaliseerde matrixvermenigvuldiging (matmul) kernel, zonder dure synchronisatie tussen thread-blokken.
• Hybrid Format: Voor training, waar het opslaan van tussenliggende activaties een bottleneck is, wordt een hybride formaat gebruikt. Dit formaat partitioneert rijen dynamisch: rijen met een laag aantal niet-nul elementen worden opgeslagen in een compact ELL-matrix, terwijl rijen met veel activaties (overflows) worden omgeleid naar een dichte back-up matrix. Dit trivialisert de opslagkosten van intermediate activaties.

2. Geoptimaliseerde CUDA Kernels:

• Inferentie: De auteurs hebben een gefuseerde kernel ontwikkeld die de "up" en "down" projecties in één enkele uitvoeringsstroom combineert. Door de gate-activaties in TwELL-formaat te gebruiken, kunnen ze de berekening van de feedforward-blokken uitvoeren met minimale globale geheugentoegang. De kernel verwerkt tegels en niet-nul elementen efficiënt, waarbij de latentie van hoog-geactiveerde tokens wordt verborgen door maximale parallelisatie.
• Training: Voor training worden custom kernels gebruikt die een hybride-naar-dense en dense-naar-hybride matrixvermenigvuldiging uitvoeren. Dit stelt de modellen in staat om backpropagation uit te voeren zonder dure dichte berekeningen, zelfs bij zeer ongelijke sparsity patronen.

3. Trainingsstrategie:

• Om hoge niveaus van sparsity te induceren zonder prestatieverlies, gebruiken de auteurs een eenvoudige L1-regularisatie op de activaties van de feedforward-lagen.
• Ze gebruiken ReLU als activatiefunctie (in plaats van SiLU), omdat ReLU inherent sparsity genereert en beter werkt met hun kernels.
• De regularisatiecoëfficiënt wordt zorgvuldig gekozen om een balans te vinden tussen sparsity en downstream prestaties.

Belangrijkste Bijdragen

1. Nieuwe CUDA Kernels: Introductie van inference- en training-kernels die ongestructureerde sparsity efficiënt exploiteren op moderne NVIDIA GPU's (H100), waardoor sparse LLMs goedkoper, sneller en lichter worden.
2. Kwantitatieve Analyse: Een studie die aantoont dat milde L1-regularisatie sparsity-niveaus van meer dan 99% kan bereiken met verwaarloosbaar verlies in downstream prestaties.
3. Schalingsvoordelen: Demonstratie dat de voordelen van deze kernels (throughput, energie, geheugen) toenemen naarmate het model groter wordt (van 0.5B tot 2B parameters).

Resultaten

De experimenten werden uitgevoerd op modellen van 0.5B tot 2B parameters, getraind op de FineWeb-dataset.

• Prestaties: Bij een L1-coëfficiënt van $2 \times 10^{-5}$ behouden de sparse modellen hun prestaties op downstream taken (zoals ARC, HellaSwag, PIQA) bijna volledig ten opzichte van dichte baselines.
• Snelheid (Inferentie): De sparse kernels leveren tot 20,5% hogere doorvoer (tokens/ms) voor inferentie op grote modellen (2B parameters).
• Snelheid (Training): Training wordt tot 21,9% sneller uitgevoerd.
• Energie en Geheugen: Er is een aanzienlijke energiebesparing (tot 17% minder energie per token) en een reductie van het piekgeheugengebruik tijdens training met meer dan 24%.
• Sparsity Patronen: De analyse toont aan dat sparsity niet uniform is; de eerste lagen zijn het minst actief, terwijl er een piek is in de vroege-midden lagen. De kernels profiteren van deze natuurlijke ongelijkheid in de data.

Betekenis en Toekomst

Dit werk is significant omdat het de "paradox" doorbreekt dat sparse berekeningen op GPU's vaak trager zijn dan dichte berekeningen. Door een nieuw formaat (TwELL) en geavanceerde kernel-fusie te introduceren, maken de auteurs sparsity tot een praktisch en schaalbaar hulpmiddel voor het verbeteren van de efficiëntie van foundation modellen.

De bevindingen suggereren dat:

• Sparsity een nieuwe as kan worden voor het schalen van LLMs, waarbij grotere modellen profiteren van grotere efficiëntiewinsten.
• Bestaande trainingsrecepten slechts minimale aanpassingen nodig hebben (L1-regularisatie + ReLU) om deze voordelen te behalen.
• De technologie ook potentieel heeft voor minder krachtige hardware (zoals RTX 6000), waar de relatieve snelheidswinst zelfs groter is dan op high-end H100 GPU's.

De auteurs hebben alle code en kernels open-source beschikbaar gesteld om de adoptie en verder onderzoek naar sparsity als een standaard voor efficiëntie in AI te versnellen.