AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

Dit artikel introduceert Activation-aware Weight Quantization (AWQ), een hardwarevriendelijke methode die de meest kritieke 1% van de LLM-weights identificeert en beschermt op basis van activatiestatistieken om efficiënte 4-bit inferentie op het apparaat mogelijk te maken, samen met het TinyChat-framework dat een snelheidswinst van meer dan 3x bereikt en de implementatie van enorme modellen zoals Llama-2 70B op mobiele GPU's mogelijk maakt.

De kern

Het Grote Probleem: De Reiskoffer

Stel je voor dat je een briljante, wereldklasse kok (een Large Language Model of LLM) hebt die verhalen kan schrijven, wiskundeproblemen kan oplossen en met je kan chatten. Deze kok is zo getalenteerd dat zijn receptenboek (het model) enorm is – ongeveer zo groot als een 350GB harde schijf.

Als je deze kok wilt meenemen op reis naar een afgelegen hut (je telefoon, laptop of auto) om te koken zonder internet, heb je een probleem: De hut is te klein om het receptenboek te bevatten. Zelfs de grootste koffers (modern computergeheugen) passen het niet. Bovendien maakt het dragen van zo'n zwaar boek dat de kok zeer traag beweegt.

Om dit op te lossen, probeerden mensen het receptenboek te verkleinen door de recepten in kleiner handschrift te schrijven (quantisatie). Maar als je alles gewoon evenredig verkleint, vergeet de kok de belangrijkste ingrediënten en smaakt het eten vreselijk.

De Oplossing: AWQ (Het Inzicht in "Uitgesproken Gewichten")

De auteurs van dit paper, het team van Ji Lin en Song Han, ontdekten een geheim: Niet alle woorden in het receptenboek zijn even belangrijk.

Stel je het receptenboek voor als een bibliotheek.

• 99% van de boeken zijn slechts naslagwerken of vulling. Deze kun je verkleinen tot tiny, 4-bits notities zonder veel smaak te verliezen.
• 1% van de boeken zijn de "Meesterrecepten". Deze bevatten de cruciale geheimen die het gerecht heerlijk maken. Als je deze verkleint, faalt de kok.

De Ontdekking: De auteurs ontdekten dat als je slechts 1% van deze "Meesterrecepten" beschermt en ze in hun originele, hoogwaardige formaat houdt, de prestaties van de kok bijna perfect blijven.

De Truc: Hoe vind je de "Meesterrecepten"?

Hier komt het slimme deel. Hoe weet je welke 1% van de boeken de "Meesterrecepten" zijn?

• Oude Manier: Je kijkt naar de boeken en raadt af welke belangrijk zijn op basis van hoe dik ze zijn (de grootte van het gewicht). Dit is als gokken dat een boek belangrijk is alleen omdat het een zware kaft heeft. Dit werkt niet goed.
• De AWQ-Manier: Je kijkt toe hoe de kok kookt. Je ziet welke boeken de kok daadwerkelijk opent en het vaakst gebruikt tijdens het maken van een gerecht (de activatie).
  • Als de kok een specifiek boek 100 keer pakt om een taart te maken, is dat boek "uitgesproken" (belangrijk).
  • AWQ zegt: "Laten we de boeken beschermen die de kok daadwerkelijk gebruikt."

De Magische Move: "Op Schaal Brengen"

Zodra ze de belangrijke boeken hebben geïdentificeerd, houden ze ze niet als enorme, zware volumes (wat alles zou vertragen). In plaats daarvan gebruiken ze een wiskundige truc genaamd Scaling.

Stel je voor dat de belangrijke boeken op een klein stukje papier staan. Om ze makkelijker leesbaar te maken (minder foutgevoelig), vergroten ze de tekst op die specifieke pagina voordat ze het hele boek verkleinen.

• Ze maken de "belangrijke" getallen iets groter.
• Hierdoor wordt de "ruis" (fouten) van het verkleinen van het boek minder opvallend voor die cruciale getallen.
• Het is als het volume opdraaien van de belangrijkste instrumenten in een orkest, zodat ze niet worden overschreeuwd wanneer het hele orkest zachter wordt.

Waarom is dit geweldig?

1. Geen Hertraining: Ze hoeven de kok niet opnieuw te leren (geen backpropagation). Ze kijken gewoon naar een paar voorbeeldgerechten (een kleine "kalibratiestel") om te zien wat de kok gebruikt.
2. Geen Overfitting: Omdat ze de voorbeeldgerechten niet uit het hoofd leren, kan de kok nog steeds heerlijke maaltijden koken voor elke keuken (coderen, wiskunde, verschillende talen) zonder in de war te raken.
3. Hardware-Vriendelijk: Ze hebben geen speciale "gemengde" koffer nodig (sommige groot, sommige klein). Ze verkleinen het hele boek, maar de "vergrootte" belangrijke delen overleven de verkleining perfect.

De Motor: TinyChat

Weten hoe je het boek moet verkleinen is één ding; het daadwerkelijk snel draaien op een klein apparaat is iets anders. De auteurs bouwden een nieuwe motor genaamd TinyChat.

Stel je TinyChat voor als een super-efficiënte leveringsvrachtwagen die specifiek is ontworpen voor deze verkleinde boeken.

• Oude Vrachtwagens: Moesten stoppen en de boeken uitpakken, lezen, verkleinen en ze dan weer inpakken elke keer dat ze verplaatsten. Zeer traag.
• TinyChat: Pakt de boeken uit terwijl het rijdt. Het voegt het uitpakken en het koken samen tot één vloeiende beweging.
• Resultaat: Op een standaard laptop of een kleine mobiele chip (zoals in een Jetson of een telefoon), draait TinyChat de verkleinde modellen 3 tot 4 keer sneller dan de standaard, niet-geoptimaliseerde versies.

De Wereldwijde Overwinningen

Het paper toont aan dat met AWQ en TinyChat:

• Je een enorm 70-miljard parameter model (zoals Llama-2-70B) kunt draaien op een enkel mobiel apparaat met 64GB geheugen, wat voorheen onmogelijk was.
• Je een 13-miljard parameter model kunt draaien op een laptop met slechts 8GB geheugen met een snelheid van 30 woorden per seconde (snel genoeg voor een real-time gesprek).
• Het werkt niet alleen voor tekst, maar ook voor multi-modale modellen (modellen die beelden zien en tekst lezen), zoals OpenFlamingo en LLaVA, zonder hun vermogen om afbeeldingen te begrijpen te verliezen.

Samenvatting

AWQ is een methode die zegt: "Verklein niet de hele hersenen gelijk. Zoek de 1% van de neuronen die het meest vuren, geef ze een kleine boost, en verklein dan de rest."
TinyChat is de software die ervoor zorgt dat deze verkleinde hersenen snel draaien op je telefoon of laptop.

Samen stellen ze ons in staat om 's werelds slimste AI-modellen uit de cloud te halen en direct in onze zakken te stoppen, geld besparen, privacy beschermen en werken zelfs wanneer het internet uitvalt.

Technische samenvatting

1. Probleemstelling

Grote Taalmodellen (LLM's) staan voor aanzienlijke uitdagingen bij implementatie op randapparaten vanwege hun enorme omvang (bijvoorbeeld vereist GPT-3 350GB in FP16) en beperkte hardwarebronnen (geheugen en rekenkracht). Hoewel Quantization-Aware Training (QAT) effectief is, is het rekenkundig duur en moeilijk te schalen. Post-Training Quantization (PTQ) is het preferente alternatief, maar bestaande methoden met lage bit-breedte (zoals GPTQ) lijden onder:

• Afname van Nauwkeurigheid: Significante prestatieverlies bij quantisatie naar zeer lage bit-breedtes (bijvoorbeeld 4-bit of 3-bit).
• Overfitting: Methoden die vertrouwen op reconstructie of backpropagatie fit vaak te goed op de kalibratieset, waardoor ze falen in generalisatie naar domeinen buiten de verdeling of verschillende modaliteiten (bijvoorbeeld multi-modale modellen).
• Hardware-inefficiëntie: Eerdere pogingen om nauwkeurigheid te behouden door een klein percentage van de gewichten in hoge precisie te houden (mixed-precision) resulteren in hardware-inefficiënties die snelheidswinst tenietdoen.

2. Methodologie: Activation-Aware Weight Quantization (AWQ)

AWQ is een hardware-vriendelijke, alleen-gewichten quantisatiemethode die werkt zonder backpropagatie of reconstructie. Het is gebaseerd op drie kerninzichten:

A. Saliente Gewichten worden Geïdentificeerd door Activatie, niet door Gewichtsgrootte

De auteurs observeren dat niet alle gewichten even belangrijk zijn. Een klein percentage (0,1%–1%) van "saliente" gewichten is cruciaal voor modelprestaties.

• Kerninzicht: Het belang van een gewichtskanaal wordt bepaald door de grootte van zijn activaties, niet door de grootte van de gewichten zelf. Kanalen met grotere activatiegroottes verwerken belangrijkere kenmerken.
• Observatie: Het behouden van slechts 1% van deze saliente kanalen in FP16 (terwijl de rest wordt gekwantiseerd) vermindert de perplexiteit drastisch (bijvoorbeeld van 43,2 naar 13,0 in OPT-6.7B). Mixed-precision is echter hardware-inefficiënt.

B. Equivalente Transformatie via Per-Kanaal Schaling

Om de hardwarekosten van mixed-precision te vermijden, leidt AWQ wiskundig af dat het opschalen van de saliente gewichtskanalen voor quantisatie hun relatieve quantisatiefout vermindert.

• Mechanisme: Als een gewicht $w$ wordt vermenigvuldigd met een schalingsfactor $s > 1$ en de bijbehorende inputactivatie $x$ wordt gedeeld door $s$, blijft de output wiskundig equivalent ($y = wx$).
• Foutreductie: De quantisatiefout is evenredig met de quantisatiestapgrootte ($\Delta$). Door saliente gewichten op te schalen, worden hun waarden groter ten opzichte van $\Delta$, wat effectief de afrondingsfout voor deze kritieke kanalen vermindert.
• Optimalisatie: Het systeem zoekt automatisch naar een optimale schalingsfactor $\alpha$ (waarbij $s = s_X^\alpha$, en $s_X$ de gemiddelde activatiegrootte is) om het outputverschil tussen het originele en het gekwantiseerde model te minimaliseren. Deze zoektocht wordt uitgevoerd via een snelle grid search over een kleine kalibratieset.

C. Data-efficiëntie en Generalisatie

• Geen Backpropagatie: AWQ vereist geen gradiëntafstijging of reconstructie, waardoor het robuust is tegen overfitting.
• Kleine Kalibratieset: Het vereist alleen het meten van de gemiddelde activatiegrootte per kanaal, waardoor het goed generaliseert naar instructie-gefine-tuned modellen en multi-modale modellen zonder domeinspecifieke fine-tuning nodig te hebben.

3. Systeemimplementatie: TinyChat

Om de theoretische geheugenbesparingen van 4-bit quantisatie om te zetten in daadwerkelijke inferentiesnelheidswinst, ontwikkelden de auteurs TinyChat, een efficiënt inferentiekader.

• Dequantisatie op het Moment: In plaats van gedequantiseerde gewichten in DRAM op te slaan (wat bandbreedte verspillen), fuseert TinyChat de dequantisatielogica direct in de matrixvermenigvuldigingskernel.
• SIMD-bewuste Gewichtsverpakking: Om te optimaliseren voor CPU/GPU SIMD-architecturen (bijvoorbeeld ARM NEON, CUDA), worden gewichten offline herschikt en verpakt. Dit maakt runtime-uitpakken mogelijk met minimale bitbewerkingen (AND, shift), wat de instructie-overhead aanzienlijk vermindert.
• Kernel-fusie: Het kader fuseert lagenormalisatie, QKV-projecties en berekeningen van positionele embedding om de overhead bij het starten van kernels en de toegang tot intermediair geheugen te minimaliseren.

4. Belangrijkste Resultaten

AWQ en TinyChat werden geëvalueerd over diverse modellen (LLaMA, OPT, Mistral, Mixtral, Vicuna, OpenFlamingo) en taken.

• Quantisatie-nauwkeurigheid:
  • AWQ presteert consistent beter dan Round-to-Nearest (RTN) en GPTQ (met en zonder herschikking) over modellen van 7B tot 70B.
  • Instructie-gefine-tuned Modellen: Bereikt bijna verliesloze prestaties op Vicuna (7B/13B) in vergelijking met FP16-baselines.
  • Multi-modale Modellen: Quantiseert succesvol OpenFlamingo en VILA-modellen, en bereikt verliesloze prestaties op 11 visueel-taalkundige benchmarks (een primeur voor low-bit VLM-quantisatie).
  • Complexe Taken: Presteert beter dan baselines op codering (MBPP) en wiskunde (GSM8K) taken, en komt in sommige 4-bit-configuraties overeen met FP16-prestaties.
• Generalisatie:
  • AWQ is robuust tegen verdelingsverschuivingen in de kalibratieset. Bij testen op verschillende datasets (bijvoorbeeld kalibreren op PubMed, evalueren op Enron) was de degradatie in perplexiteit van AWQ minimaal (0,5–0,6) in vergelijking met GPTQ (2,3–4,9).
  • Het vereist een kalibratieset 10x kleiner dan GPTQ om vergelijkbare prestaties te bereiken.
• Inferentiesnelheid (TinyChat):
  • Snelheidswinst: Bereikt 3,2× tot 3,9× snelheidswinst ten opzichte van HuggingFace FP16-implementaties op desktop (RTX 4090) en mobiele GPU's (Jetson Orin).
  • Implementatie: Maakt de implementatie van Llama-2-70B mogelijk op een enkele Jetson Orin (64GB RAM) en Llama-2-13B op een laptop met slechts 8GB RAM (33 tokens/sec), wat onmogelijk is met FP16.
  • Randapparaten: Voert 7B-modellen uit op een Raspberry Pi 4B met 0,7 tokens/sec.

5. Betekenis en Impact

• Democratisering van LLM's: AWQ en TinyChat maken het haalbaar om state-of-the-art LLM's (inclusief modellen met 70B parameters) uit te voeren op consumentenhardware, mobiele apparaten en IoT-randknooppunten, waardoor de afhankelijkheid van cloudinfrastructuur afneemt.
• Privacy en Kosten: Door lokale uitvoering mogelijk te maken, wordt de privacy van gebruikers versterkt en worden cloudlatentie en kosten geëlimineerd.
• Generalisatie: In tegenstelling tot eerdere methoden die moeite hebben met instructie-gefine-tuned of multi-modale modellen, behoudt AWQ het "generalist"-karakter van LLM's, waardoor het een veelzijdige oplossing is voor diverse AI-toepassingen.
• Adoptie: De methode is breed overgenomen door grote spelers in de industrie en open-source projecten, waaronder HuggingFace Transformers, NVIDIA TensorRT-LLM, Microsoft DirectML en vLLM.

Kortom, AWQ biedt een wiskundig onderbouwde, hardware-efficiënte oplossing voor low-bit LLM-quantisatie, terwijl TinyChat ervoor zorgt dat deze theoretische winsten worden gerealiseerd als praktische, hoge snelheid inferentie op randapparaten.