Onderzoekers hebben een nieuwe optimizer ontwikkeld, genaamd SkewAdam, die het geheugengebruik tijdens het trainen van Mixture-of-Experts (MoE) taalmodellen drastisch vermindert zonder in te boeten op nauwkeurigheid. De methode maakt gebruik van een gelaagde toewijzing van de optimizer-status, waardoor grote modellen op hardware met beperkter geheugen kunnen worden getraind.
Het trainen van moderne taalmodellen vereist enorme hoeveelheden computergeheugen, waarbij vooral de zogenaamde 'optimizer state' een grote hap uit het budget neemt. In een recent onderzoek gepubliceerd via arxiv.org legt auteur Nuemaan Malik uit dat bij een MoE-model met 6,78 miljard parameters de veelgebruikte AdamW-optimizer maar liefst 50,6 GB aan geheugen reserveert voor het bijhouden van eerste en tweede momenten. Dit staat in schril contrast met de eigenlijke gewichten van het model, die in bfloat16-formaat slechts 12,6 GB beslaan.
Een gelaagde aanpak voor geheugenefficiëntie
De kern van de nieuwe oplossing, SkewAdam, is de observatie dat verschillende onderdelen van een MoE-model verschillende statistische eigenschappen hebben. In plaats van voor elke parameter dezelfde hoeveelheid geheugen te reserveren, introduceert SkewAdam een gelaagde toewijzing (tiered allocation) gebaseerd op drie specifieke populaties binnen het model:
- De dense backbone: Deze beslaat ongeveer 5% van de parameters. Voor dit deel behoudt SkewAdam de float32-momentum en een gefactoreerd tweede moment.
- De experts: Dit is het grootste deel van het model (circa 95%). Hier wordt enkel een gefactoreerd tweede moment bijgehouden.
- De router: Dit minuscule deel (minder dan 0,01%) krijgt een exacte weergave van het tweede moment.
Volgens de details in de resulteert deze strategie in een optimizer-status van slechts 1,29 GB. Dit is ongeveer 2,6% van het geheugen dat AdamW zou vereisen. In de praktijk leidde dit ertoe dat het piekgeheugengebruik tijdens de training daalde van 81,4 GB naar 31,3 GB. Hierdoor past het trainingsproces binnen het budget van een accelerator met 40 GB geheugen, wat de toegankelijkheid van het trainen van grote modellen aanzienlijk vergroot.
Prestaties en vergelijkingen
Om de effectiviteit van SkewAdam te testen, is het model vergeleken met andere bekende optimizers zoals AdamW, Muon en Lion over een dataset van 82 miljoen tokens. De resultaten waren opvallend: SkewAdam behaalde een validatie-perplexiteit van 108,4. Ter vergelijking: AdamW scoorde 126,8, Muon 120,2 en Lion 393,7. Een lagere perplexiteit duidt op een beter presterend model.
Het onderzoek benadrukt dat de geheugenbesparing zelf niet direct verantwoordelijk is voor de hogere nauwkeurigheid. Uit een zogenaamde 'tier ablation' bleek dat een configuratie met twintig keer meer statusgeheugen dezelfde resultaten behaalde. De winst in nauwkeurigheid komt voort uit het behoud van momentum, iets wat optimizers zoals Adafactor laten vallen. Adafactor bleek in de tests namelijk 40 punten achter te blijven op SkewAdam.
Conclusie over modelarchitectuur
Zelfs na het optimaliseren van de leer-rates (learning rates) van de concurrerende methoden bleef SkewAdam superieur. De best afgestelde versie van AdamW bereikte een perplexiteit van 118,5, terwijl een getunede Adafactor op 139,7 uitkwam.
De conclusie van de auteur, zoals beschreven in de , is dat de locatie en de wijze waarop de optimizer-status wordt opgeslagen, minstens zo belangrijk zijn als de totale hoeveelheid geheugen die beschikbaar is. Door specifiek in te spelen op de architectuur van Mixture-of-Experts modellen, kan SkewAdam een efficiënter alternatief bieden voor de huidige industriestandaarden.
De volledige resultaten, inclusief trainingslogs en de bijbehorende code, zijn beschikbaar gesteld via de van het onderzoek.