Technische Samenvatting: PagedWeight

Probleemstelling

Mixture-of-Experts (MoE) Large Language Models (LLMs) bieden hoge efficiëntie en nauwkeurigheid door slechts een subset van experts per token te activeren. Echter, in scenario's voor modelbediening (serving) worden MoE-modellen geconfronteerd met een kritieke geheugenspanning: de GPU moet zowel de enorme modelgewichten als de groeiende Key-Value (KV) cache opslaan. Hoewel technieken zoals PagedAttention fragmentatie van de KV-cache beheren, pakken zij de significante geheugenvoetafdruk van MoE-gewichten niet aan, die meer dan 60% van het GPU-geheugen kunnen innemen.

Bestaande kwantiseringsmethoden zijn grotendeels statisch; zij fixeren de precisie voordat de inferentie begint. Hoewel er enkele runtime precisie-adaptatiemethoden bestaan, maken deze geen gebruik van precisieveranderingen om GPU-geheugen dynamisch te heralloceren tussen gewichten en de KV-cache. Bijgevolg is er een gebrek aan systemen die in staat zijn om dynamisch geheugen vrij te maken van minder kritieke expert-gewichten naarmate de KV-cache groeit, zonder de taaknauwkeurigheid op te offeren of de inferentie te onderbreken.

Methodologie: PagedWeight

De auteurs stellen PagedWeight voor, een nieuw beheersysteem voor de bediening van MoE LLM's dat tijdens runtime de modelgewichten kwantiseert om de balans tussen de precisie van expert-gewichten en de grootte van de KV-cache te beheren. Het systeem behandelt expert-gewichten vergelijkbaar met gepageneerde KV-cacheblokken, waardoor dynamische offloading en herladen mogelijk is.

Kerncomponenten

1. Paged Weight Representatie:

   • PagedWeight werkt op de granulariteit van per-expert lineaire blokken (geïdentificeerd door laag, expert en module type: gate_up of down).
   • Het maakt gebruik van Any-Precision LLM (APL) technologie, waarbij gewichten worden gerepresenteerd in een overlappend bit-plane formaat met bijbehorende Lookup Tables (LUTs).
   • Een Weight Page Table houdt de gewenste bitbreedte ($d_i$) en de gecommitteerde bitbreedte ($q_i$) bij voor elk lineair blok. Pagina's kunnen zich in een GPU-resident, CPU-resident of transfer-status bevinden.

2. Kwaliteitsbewuste Runtime Planner:
   De planner bepaalt welke gewichten gekwantiseerd moeten worden (de bitbreedte verlagen) op basis van drie factoren om nauwkeurigheidsverlies te minimaliseren:

   • Offline Globale Sensitiviteit: Gebruikt Hessian-gewogen scores om de intrinsieke sensitiviteit van elk lineair blok voor kwantisering te schatten.
   • Online Routing Statistieken: Houdt de "routing mass" (frequentie en gewicht van selectie) van experts bij. Veel gerouteerde ("hot") experts worden beschermd met hogere bitbreedte-vloeren en schade-multiplicatoren.
   • Prompt Residuals: Een prompt-specifieke correctieterm die voorspeld wordt via lineaire regressie-koppen. Dit past de globale sensitiviteitschat aan op basis van de huidige inputsequentie, in de wetenschap dat de impact van kwantisering varieert per prompt.

   De planner berekent een "schade"-score voor potentiële reducties in bitbreedte. Wanneer de KV-cache druk een geheugendoel triggert, selecteert de planner gulzig de laagste-schade reducties (hoogste kwaliteit per bespaarde byte) om het doel te bereiken.

3. Asynchrone Executie & Kernel Optimalisatie:

   • Asynchrone Pagina Beweging: Om latentie te verbergen, verplaatst het systeem gewichtspagina's naar het CPU-RAM en laadt deze asynchroon opnieuw in. Bitbreedte-verplichtingen worden pas bij veilige grenzen bijgewerkt (bijv. nadat een plan volledig is uitgevoerd), zodat de inferentie niet wordt onderbroken.
   • Fused Mixed-Precision Kernel: Een aangepaste CUDA-kernel leest direct Any-Precision bit-planes en LUTs, waarbij verschillende bitbreedtes voor elk lineair blok binnen één gefuseerde operatie worden ondersteund om overhead te verminderen.

Belangrijkste Bijdragen

1. Systeemontwerp: De voorstel van PagedWeight, een gepageneerd gewichtssysteem voor online Any-Precision MoE-gewichten dat gecommitteerde/gewenste bitbreedtes, paginastatus en geheugengebruik dynamisch beheert.
2. Planningsalgoritme: Een kwaliteitsbewuste runtime planner die offline sensitiviteit, online routing statistieken en prompt residuals synthetiseert om strategieën voor lage-schade gewichtsreductie te selecteren onder geheugendruk.
3. Executiestrategie: Implementatie van asynchrone pagina-beweging om de offload/reload latentie te verbergen met minimale verlies van throughput, gekoppeld aan een gefuseerde mixed-precision MoE kernel.
4. Empirische Validatie: Uitgebreide evaluatie over drie MoE-modellen (Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B) die superieure prestaties aantonen ten opzichte van bestaande baselines.

Resultaten

De auteurs hebben PagedWeight geëvalueerd tegenover FP16, Any-Precision LLM (APL), DP-LLM en MxMoE over taken op het gebied van taalmodellering, redeneren en lange context.

• Kwaliteit-Geheugen Tradeoff: PagedWeight bereikt FP16-equivalente nauwkeurigheid met tot wel 72,0% GPU-geheugenbesparing en een 1,94× verbetering in throughput vergeleken met baselines.
• Kwaliteit bij Vast Budget: Bij vergelijkbare geheugenbudgetten verbetert PagedWeight de taakkwaliteit ten opzichte van kwantiseringsmethoden met tot wel 39,3%, met een throughput-verlies van maximaal 4,1%.
• Long-Context Prestaties: Onder strikte geheugenbudgetten (bijv. 10 GB) behoudt PagedWeight de long-context kwaliteit (Passage Retrieval, NarrativeQA) die vergelijkbaar is met FP16, terwijl statische kwantiseringsbaselines (zoals APL) significante degradatie vertonen.
• Throughput: PagedWeight komt overeen met de throughput van uniforme kwantiseringsbaselines, met de grootste waargenomen daling van 4,1% bij een batch size van 4.
• Ablatie: Het verwijderen van componenten zoals routing statistieken, prompt residuals of pagina-beweging verslechtert consequent de prestaties, wat de noodzaak van het volledige systeemontwerp bevestigt.

Betekenis

Het artikel claimt dat PagedWeight effectief navigeert door de complexe afweging tussen modeltaak-nauwkeurigheid, geheugengebruik en throughput/latentie. Door expert-gewichten te behandelen als gepageneerde, beheerbare eenheden in plaats van statische activa, ontsluit PagedWeight een nog onverkende tradeoff-ruimte tussen taaknauwkeurigheid en de allocatie van GPU-geheugen tussen gewichten en de KV-cache. Het systeem demonstreert dat dynamische, kwaliteitsbewuste kwantisering een levensvatbare en superieure strategie is voor MoE-bediening vergeleken met statische kwantisering of methoden die niet adapteren aan runtime geheugendruk.