Technisches Resümee: PagedWeight

Problemstellung

Mixture-of-Experts (MoE) Large Language Models (LLMs) bieten hohe Effizienz und Genauigkeit, indem sie pro Token nur eine Teilmenge von Experten aktivieren. In Serving-Szenarien stehen MoE-Modelle jedoch vor einem kritischen Speicherkonflikt: Die GPU muss sowohl die massiven Modellgewichte als auch den wachsenden Key-Value (KV)-Cache speichern. Während Techniken wie PagedAttention die Fragmentierung des KV-Caches bewältigen, adressieren sie nicht den erheblichen Speicherbedarf der MoE-Gewichte, die über 60 % des GPU-Speichers einnehmen können.

Bestehende Quantisierungsmethoden sind weitgehend statisch und legen die Präzision vor Beginn der Inferenz fest. Obwohl einige Methoden zur Präzisionsanpassung zur Laufzeit existieren, nutzen diese die Präzisionsänderungen nicht, um den GPU-Speicher dynamisch zwischen Gewichten und dem KV-Cache neu zu verteilen. Folglich mangelt es an Systemen, die in der Lage sind, dynamisch Speicher von weniger kritischen Experten-Gewichten freizugeben, während der KV-Cache expandiert, ohne die Aufgabenpräzision zu opfern oder die Inferenz zu unterbrechen.

Methodik: PagedWeight

Die Autoren schlagen PagedWeight vor, ein neuartiges Verwaltungssystem für das Serving von MoE-LLMs, das Modellgewichte zur Laufzeit dynamisch quantisiert, um die Präzision der Experten-Gewichte mit der Größe des KV-Caches in Einklang zu bringen. Das System behandelt Experten-Gewichte ähnlich wie gepagte KV-Cache-Blöcke, was ein dynamisches Offloading und Reloading ermöglicht.

Kernkomponenten

1. Paged-Weight-Repräsentation:

   • PagedWeight operiert auf der Granularität von pro-Experte lineare Blöcken (identifiziert durch Schicht, Experte und Modultyp: gate_up oder down).
   • Es nutzt Any-Precision LLM (APL)-Technologie, welche die Gewichte in einem überlagerten Bit-Plane-Format mit zugehörigen Lookup-Tabellen (LUTs) darstellt.
   • Eine Weight Page Table verfolgt die gewünschte Bitbreite ($d_i$) und die zugesagte (committed) Bitbreite ($q_i$) für jeden linearen Block. Pages können sich im Zustand im GPU-Resident, CPU-Resident oder im Transfer-Status befinden.

2. Qualitätsbewusster Runtime-Planner:
   Der Planner bestimmt, welche Gewichte zu quantisieren sind (Bitbreite reduzieren), basierend auf drei Faktoren, um den Genauigkeitsverlust zu minimieren:

   • Offline Global Sensitivity: Verwendet Hessian-gewichtete Scores, um die intrinsische Sensitivität jedes linearen Blocks gegenüber der Quantisierung zu schätzen.
   • Online Routing Statistics: Verfolgt die „Routing Mass“ (Häufigkeit und Gewicht der Auswahl) von Experten. Häufig geroutete („heiße“) Experten werden durch höhere Bitbreiten-Untergrenzen und Schadensmultiplikatoren geschützt.
   • Prompt Residuals: Ein promptspezifischer Korrekturterm, der via linearer Regressionsköpfe vorhergesagt wird. Dies passt die globale Sensitivitätsschätzung basierend auf der aktuellen Eingabesequenz an, da die Auswirkungen der Quantisierung je nach Prompt variieren.

   Der Planner berechnet einen „Damage“-Score für potenzielle Bitbreiten-Reduktionen. Wenn der KV-Cache-Druck ein Speicherziel auslöst, wählt der Planner gierig (greedy) die Reduktionen mit dem geringsten Schaden (höchste Qualität pro eingespartem Byte) aus, um das Ziel zu erreichen.

3. Asynchrone Ausführung & Kernel-Optimierung:

   • Asynchrone Page-Bewegung: Um Latenzen zu verbergen, lagert das System Weight-Pages in den CPU-RAM aus und lädt sie asynchron nach. Die Bitbreiten-Zusagen (commitments) werden erst an sicheren Grenzen aktualisiert (z. B. nachdem ein Plan vollständig ausgeführt wurde), um sicherzustellen, dass die Inferenz nicht unterbrochen wird.
   • Fused Mixed-Precision Kernel: Ein benutzerdefinierter CUDA-Kernel liest Any-Precision Bit-Planes und LUTs direkt aus und unterstützt unterschiedliche Bitbreiten für jeden linearen Block innerhalb einer einzigen gefügten Operation, um den Overhead zu reduzieren.

Wichtigste Beiträge

1. Systemdesign: Der Vorschlag von PagedWeight, einem paged-weight System für Online Any-Precision MoE-Gewichte, das committed/desired Bitbreiten, Page-Zustände und Speicherverbrauch dynamisch verwaltet.
2. Planungsalgorithmus: Ein qualitätsbewusster Runtime-Planner, der Offline-Sensitivität, Online-Routing-Statistiken und Prompt-Residuals synthetisiert, um Strategien zur Reduktion der Gewichte mit geringem Schaden unter Speicherdruck zu wählen.
3. Ausführungsstrategie: Implementierung von asynchroner Page-Bewegung, um die Offload/Reload-Latenz mit minimalem Durchsatzverlust zu verbergen, gekoppelt mit einem gefügten Mixed-Precision MoE-Kernel.
4. Empirische Validierung: Umfassende Evaluierung über drei MoE-Modelle (Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B), die eine überlegene Leistung gegenüber bestehenden Baselines demonstriert.

Ergebnisse

Die Autoren evaluierten PagedWeight gegen FP16, Any-Precision LLM (APL), DP-LLM und MxMoE bei Aufgaben der Sprachmodellierung, des logischen Schlussfolgerns (Reasoning) und bei Long-Context-Aufgaben.

• Qualitäts-Speicher-Tradeoff: PagedWeight erreicht eine FP16-äquivalente Genauigkeit mit bis zu 72,0 % GPU-Speichereinsparung und einer 1,94× Durchsatzverbesserung im Vergleich zu Baselines.
• Qualität bei festem Budget: Bei ähnlichen Speicherbudgets verbessert PagedWeight die Aufgabenqualität gegenüber Quantisierungsmethoden um bis zu 39,3 %, bei einem Durchsatzverlust von höchstens 4,1 %.
• Long-Context-Performance: Unter engen Speicherbudgets (z. B. 10 GB) behält PagedWeight eine Long-Context-Qualität (Passage Retrieval, NarrativeQA) bei, die mit FP16 vergleichbar ist, während statische Quantisierungs-Baselines (wie APL) eine signifikante Degradierung erleiden.
• Durchsatz: PagedWeight erreicht den Durchsatz von uniformen Quantisierungs-Baselines, wobei der größte beobachtete Abfall 4,1 % bei einer Batch-Größe von 4 beträgt.
• Ablation: Das Entfernen von Komponenten wie Routing-Statistiken, Prompt-Residuals oder Page-Bewegung verschlechtert die Leistung konsistent, was die Notwendigkeit des vollständigen Systemdesigns bestätigt.

Bedeutung

Das Paper behauptet, dass PagedWeight effektiv den komplexen Tradeoff zwischen Modell-Aufgabenpräzision, Speicherverbrauch und Durchsatz/Latenz navigiert. Indem es Experten-Gewichte als gepagte, verwaltbare Einheiten statt als statische Assets behandelt, erschließt PagedWeight einen bisher unentdeckten Tradeoff-Raum zwischen Aufgabenpräzision und der Allokation von GPU-Speicher zwischen Gewichten und dem KV-Cache. Das System zeigt, dass eine dynamische, qualitätsbewusste Quantisierung eine praktikable und überlegene Strategie für das MoE-Serving im Vergleich zu statischer Quantisierung oder Methoden ist, die nicht auf den Runtime-Speicherdruck reagieren.