Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling

Technische Samenvatting: Geheugenefficiënt LLM-trainen met Dynamische Sparsity

Probleemstelling
Dynamic Sparse Training (DST) biedt een paradigma voor het verbeteren van de efficiëntie van diepe neurale netwerken door een ijle (sparse) topologie tijdens de training te behouden. Het toepassen van DST op Large Language Models (LLMs) is echter problematisch gebleken vanwege optimalisatie-instabiliteit. De auteurs identificeren dat de naïeve toepassing van standaard Adam-gebaseerde optimizers op DST leidt tot ernstige "loss spikes" (verliespieken) direct na topologie-updates (pruning en regrowth).

De kernoorzaak wordt geïdentificeerd als een "cold-start"-effect: nieuw teruggegroeide (regrown) parameters missen de geaccumuleerde optimizer-states (eerste en tweede moment-schattingen). In Adam resulteren de bias-correctietermen voor deze nieuwe parameters in disproportioneel grote update-magnitudes ten opzichte van volwassen gewichten. Dit verstoort het trainingsverloop, wat leidt tot transiënte loss spikes en de convergentie belemmert, vooral wanneer de modelomvang toeneemt of de topologie-updates agressiever worden. Bovendien slagen conventionele DST-implementaties er vaak niet in om geheugenbesparingen te realiseren omdat ze volledige optimizer-states alloceren voor de gehele parameterruimte, ook al is slechts een subset van de parameters actief.

Methodologie: Sparse Memory-Efficient Training (SMET)
Om deze uitdagingen aan te pakken, stellen de auteurs SMET voor, een framework ontworpen om DST te stabiliseren en het geheugengebruik voor LLMs te verminderen. De methode bestaat uit drie kernelementen:

1. Optimizer-State Warm-up: Om het cold-start-effect te mitigeren, gebruikt SMET een tweefasige warm-up strategie voor nieuw teruggegroeide parameters:

   • Implicit Warm-up (Timestep Reset): De interne tijdstap-teller ($t_i$) voor teruggegroeide parameters wordt gereset. Dit activeert het bias-correctiemechanisme van Adam opnieuw, waardoor de noemer van de update-regel wordt opgeblazen en de initiële update-magnitude op natuurlijke wijze wordt beperkt.
   • Explicit Warm-up (Learning Rate Scaling): Een lokale lineaire warm-up wordt toegepast waarbij de learning rate voor teruggegroeide parameters lineair toeneemt van nul naar de doelrate over $W$ stappen.
   • Deze mechanismen voorkomen gezamenlijk explosieve updates en verzachten het optimalisatieverloop direct na topologieveranderingen.

2. Density-Aware Learning Rate Scaling: Om de conservatieve updates die door de warm-up worden geïntroduceerd te compenseren en rekening te houden met de verminderde effectieve fan-in van ijle lagen, schaalt SMET de learning rate. Geïnspireerd door Maximal Update Parameterization ($\mu$P), wordt de learning rate geschaald met $1/\sqrt{d}$, waarbij $d$ de huidige sparsity-densiteit is. Dit zorgt ervoor dat de update-magnitudes vergelijkbaar blijven met dense training-regimes.

3. Memory-Efficient Optimization: SMET implementeert een ijle opslagstrategie waarbij gradiënten en optimizer-states (moment-schattingen $m$ en $v$) alleen worden opgeslagen voor actieve (niet-nul) parameters.

   • Gradiënten worden uitsluitend behouden bij actieve indices tijdens backpropagation.
   • Om de overhead van het bijhouden van actieve indices en tijdstappen te beheren, wordt een block-wise opslagstrategie gebruikt, waarbij parameters worden gegroepeerd in blokken van grootte $B$ om de metadata-kosten te amortiseren.
   • Deze aanpak behoudt de sparsity gedurende de gehele trainingscyclus, in tegen tegenstelling tot methoden die states comprimeren maar dense modellen produceren.

Belangrijkste Bijdragen

1. Analyse van Instabiliteit: Het artikel identificeert en analyseert theoretisch het fenomeen van loss spikes in LLM pre-training met DST, waarbij het de oorzaak wijst aan het cold-start-effect van nieuw teruggegroeide gewichten in Adam optimizers.
2. Stabilisatie-framework: De auteurs stellen SMET voor, dat optimizer-state warm-up en density-aware learning-rate scaling integreert om topologie-updates te stabiliseren en de prestaties van de training te verbeteren.
3. Geheugenefficiënte Implementatie: Er is een ijle optimalisatie-implementatie ontwikkeld die gradiënten en optimizer-states alleen opslaat voor actieve parameters, waarbij zowel unstructured als block-wise sparsity-patronen worden ondersteund terwijl de sparsity gedurende de training behouden blijft.
4. Systematische Evaluatie: Het werk biedt een uitgebreide evaluatie van DST over diverse model-schalen en sparsity-niveaus, waarbij wordt aangetoond dat SMET bestaande DST-methoden overtreft en de prestaties van dense training benadert.

Experimentele Resultaten
Experimenten werden uitgevoerd op LLaMA-gebaseerde modellen (variërend van 60M tot 1B parameters) die werden gepre-traind op de C4 en OpenWebText datasets.

• Prestaties: SMET presteert consistent beter dan andere geheugenefficiënte baselines (LoRA, ReLoRA, GaLore) en bestaande DST-methoden (SET, RigL). Bij grotere model-schalen (bijv. 1B parameters) bereikt SMET een perplexiteit die dicht bij dense training ligt, zelfs bij hoge sparsity-niveaus (bijv. densiteit $d=0.1$).
• Stabiliteit: De methode elimineert effectief de loss spikes die worden waargenomen in vanilla DST, wat resulteert in vloeiende trainingscurves.
• Geheugenefficiëntie: Door states alleen voor actieve parameters op te slaan, vermindert SMET de geheugenvoetafdruk van gradiënten en optimizer-states aanzienlijk. Voor een 1B model bij densiteit $d=0.4$ worden substantiële geheugenbesparingen waargenomen.
• Schaalbaarheid: Het prestatieverschil tussen sparse en dense training wordt kleiner naarmate de modelgrootte toeneemt, wat suggereert dat sparse training bijzonder levensvatbaar is voor grootschalige modellen.
• Generalisatie: De voordelen van SMET generaliseren over verschillende regrowth-strategieën (random en gradient-based) en verschillende model-architecturen (LLaMA en Qwen).

Betekenis en Claims
Het artikel beweert dat SMET de weg vrijmaakt voor sparse training als een praktisch alternatief voor dense training voor LLMs. Door de optimalisatie-instabiliteit die inherent is aan dynamic sparse training op te lossen en werkelijke geheugenbesparingen te realiseren door middel van sparse state storage, maakt SMET een stabiele, schaalbare en geheugenefficiënte pre-training mogelijk. De auteurs merken op dat hoewel SMET de trainingsstabiliteit verbetert en de optimizer-geheugens bij de laag verlaagt, er nog steeds uitdagingen zijn met betrekking tot activatiegeheugen en de noodzaak voor efficiënte hardware-ondersteuning voor fine-grained sparsity om de inference-versnellingen volledig te realiseren. Desalniettemin toont dit werk aan dat met passende stabilisatie en geheugenbeheer, dynamische sparsity succesvol kan worden opgeschaald naar foundation models.