AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs

1. Probleemstelling

Het trainen van Large Language Models (LLM's) vereist enorme GPU-geheugencapaciteit, die voornamelijk wordt verbruikt door modelparameters, optimizer-toestanden, gradiënten en activaties.

• Geheugenflesnek: Activaties nemen tijdens het trainen doorgaans het grootste deel van het geheugen in, met een lineaire schaling met sequentielengte en batchgrootte.
• Beperkingen van bestaande methoden:
  • Activatieherberekening/Offloading: Vermindert het geheugengebruik maar introduceert aanzienlijke computerefficiëntie-overhead (herberekening) of I/O-latentie (offloading).
  • Standaard kwantisatie: Hoewel 8-bit kwantisatie (INT8/FP8) gebruikelijk is, leidt toepassing op 4-bit activaties of 8-bit gradiënten vaak tot ernstig verlies aan nauwkeurigheid, trage convergentie of trainingsinstabiliteit door numerieke overloop tijdens accumulatie en All-Reduce-operaties.
  • Gradiëntcommunicatie: Huidige systemen slaan gradiënten vaak op in hoge precisie (FP32) en kwantiseren ze pas voor communicatie, waardoor het geheugengebruik van opgeslagen gradiënten niet wordt verminderd.

2. Methodologie: AGoQ Framework

De auteurs stellen AGoQ voor, een holistisch framework dat twee nieuwe technieken integreert om opslag van activaties tot bijna 4-bit en opslag/communicatie van gradiënten tot 8-bit te realiseren zonder in te leveren op modelconvergentie.

A. Layer-Aware Activation Quantization (LAAQ)

In plaats van alle lagen uniform te kwantiseren naar 4 bits (wat nauwkeurigheidsverlies veroorzaakt), wijst LAAQ dynamisch bit-breedtes toe op basis van laagtypes en pijplijnstadia.

• Theoretische foutanalyse: De auteurs leiden foutgrenzen af voor verschillende operaties (RMSNorm, SiLU, Attention, GEMM).
  • Vondsten: Het kwantiseren van activaties van de Attention-module (Q, K, V) introduceert aanzienlijk grotere gradiëntfouten in vergelijking met andere modules. Daarom worden Attention-activaties niet gekwantiseerd (behouden in hogere precisie).
  • Strategie: Voor andere lagen (RMSNorm, SiLU, FFN) gebruikt het systeem 4-bit block-wise kwantisatie.
  • Herberekening versus Caching: De analyse toont aan dat voor operaties zoals RMSNorm en SiLU, het opslaan van alleen de gekwantiseerde invoer en het herberekenen van tussentijdse waarden tijdens de backward-pass een strakkere foutgrens oplevert dan het cacheën van tussentijdse waarden in gekwantiseerde vorm. AGoQ slaat daarom alleen 4-bit invoer op en berekent intermediairen opnieuw.
• Dynamic Bit-width Compensation (DBCA-PP):
  • Bij Pipeline Parallelism (PP) bevinden verschillende pijplijnstadia zich in verschillende aantallen activatie-batches (ongelijke geheugengebruik).
  • AGoQ maakt hier gebruik van door hogere bit-breedtes (bijv. 5, 6 of 8 bits) toe te wijzen aan stadia met minder actieve batches en lagere bit-breedtes (4 bits) aan stadia met meer batches. Dit balanceert het totale geheugengebruik terwijl een gemiddelde van ~4 bits per element wordt behouden, wat compensatie biedt voor precisieverlies waar geheugen beschikbaar is.

B. Precision-Preserved Gradient Quantization (QuanGrad)

Deze techniek vermindert zowel het geheugengebruik als de communicatietijd door 8-bit (FP8) gradiënten gedurende het hele proces te gebruiken.

• Lokale accumulatie: Om overloop te voorkomen bij het accumuleren van gradiënten van meerdere mini-batches, dekwantiseert het systeem de FP8 hoofdgradiënt naar FP16/BF16, voegt de nieuwe lokale gradiënt toe en kwantiseert de som vervolgens terug naar FP8.
• Precision-Preserved All-Reduce: Standaard FP8 All-Reduce veroorzaakt overloop tijdens de reduceringsstap. AGoQ splitst de All-Reduce-operatie op in:
  1. All-to-All: Verstuur gecomprimeerde FP8-gradiënten naar alle apparaten.
  2. Lokale reduce: Elk apparaat dekwantiseert ontvangen data naar FP32, voert de lokale som uit en kwantiseert terug naar FP8.
  3. All-Gather: Verzamel de gereduceerde resultaten.
  • Hierdoor gebeurt de accumulatie in hoge precisie (FP32) terwijl de communicatie in lage precisie (FP8) blijft, wat overloop voorkomt en convergentie behoudt.

C. Kernel Fusion

Om de overhead van kwantisatie/dekwantisatie-operaties te elimineren, fuseert AGoQ deze element-wise operaties met aangrenzende GEMM-kernen (General Matrix Multiply). Dit maakt gebruik van Tensor Cores voor het zware werk, terwijl de kwantisatielogica binnen dezelfde kernel op CUDA-kernen blijft, wat de latentie minimaliseert.

3. Belangrijkste Bijdragen

1. Layer-Aware Quantization Strategy: Een theoretische en praktische aanpak om te identificeren welke lagen veilig naar 4 bits kunnen worden gecomprimeerd (exclusief Attention) en welke strategieën voor herberekening vereisen om fouten te minimaliseren.
2. Dynamic Bit-width Compensation: Een nieuw algoritme voor Pipeline Parallelism dat de kwantisatieprecisie dynamisch aanpast op basis van geheugenbeschikbaarheid in elk stadium, waardoor geheugenefficiëntie wordt gemaximaliseerd zonder piekgebruik te verhogen.
3. Precision-Preserved Gradient Communication: Een methode om 8-bit All-Reduce uit te voeren door de operatie te splitsen om accumulatie in hoge precisie mogelijk te maken, waardoor de overloopproblemen worden opgelost die eerder trainen met lage-bit gradiënten verhinderden.
4. System Integration: Volledige integratie in Megatron-LM, compatibel met 8-bit Adam-optimizers, met ondersteuning voor zowel NVIDIA GPU's als Huawei Ascend NPUs.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op clusters met tot 64 GPU's (NVIDIA A6000/Pro 6000) met LLaMA-modellen (7B, 8B, 13B, 34B) en OLMo-1B.

• Geheugenreductie:

  • AGoQ vermindert het piekgeheugengebruik met maximaal 52% in vergelijking met training in volledige precisie (Megatron-LM).
  • In vergelijking met het state-of-the-art COAT-systeem vermindert AGoQ het activatiegeheugen met een extra 30% en het gradiëntgeheugen met 75%.
  • Het maakt het trainen van langere sequentielengtes mogelijk (bijv. 80K tokens) die in basissystemen leiden tot Out-Of-Memory (OOM)-fouten.

• Trainingsnelheid:

  • Bereikt tot 1,34× versnelling in end-to-end trainingsduur in vergelijking met Megatron-LM (met of zonder ZeRO) en DeepSpeed.
  • De versnelling is duidelijker bij langere sequentielengtes door verminderde communicatie-overhead en geheugenbeperkingen.

• Convergentie en Nauwkeurigheid:

  • Convergentie: Trainingsverliescurves op LLaMA2-7B en LLaMA3-8B komen nauw overeen met de baseline in volledige precisie. In tegenstelling hiermee vertoonden naïeve FP8 All-Reduce-methoden aanzienlijk hoger verlies.
  • Downstream Tasks: Zero-shot nauwkeurigheid op datasets (ARC, HellaSwag, PIQA, enz.) blijft vergelijkbaar met de baseline, waarbij sommige taken lichte verbeteringen tonen.
  • Stabiliteit: Gradiëntnormen blijven stabiel gedurende de training, wat de effectiviteit van de precisiebehoudende accumulatie bevestigt.

5. Betekenis

• Mogelijk maken van Long-Context Training: Door het geheugengebruik drastisch te verminderen, maakt AGoQ het haalbaar om LLM's met zeer lange contextvensters (bijv. 64K–80K tokens) te trainen op bestaande hardware zonder enorme clusters te vereisen.
• Het gat dichten: Het duwt kwantisatie succesvol naar 4-bit voor activaties en 8-bit voor gradiënten, een regime dat eerder als te instabiel werd beschouwd voor LLM-training, zonder in te leveren op modelkwaliteit.
• Schaalbaarheid: De methode is zeer schaalbaar over verschillende parallelisme-strategieën (DP, TP, PP) en hardware-architecturen, en biedt een praktische weg naar het trainen van modellen met biljoenen parameters op beperkte middelen.

Kortom, AGoQ vertegenwoordigt een aanzienlijke vooruitgang in systemen voor LLM-training, die verder gaat dan eenvoudige compressie naar een wiskundig onderbouwde, adaptieve kwantisatieframework dat geheugenefficiëntie, communicatiesnelheid en modelnauwkeurigheid in evenwicht brengt.