FlashAttention-Skalierung, Multi-GPU-Training und Gewichtsquantisierung

Moderne KI-Infrastruktur nutzt hardwareoptimierte Attention-Kernels, verteilten Multi-GPU-Parallelismus und Quantisierung mit geringer Bitbreite, um riesige Modelle effizient zu trainieren und bereitzustellen. Diese drei Techniken überbrücken die Lücke zwischen theoretischer Modellgröße und den physischen Grenzen von GPUs.

1. FlashAttention-2 und FlashAttention-3 auf Hopper-GPUs (H100)

Während das ursprüngliche FlashAttention den Engpass der Speicherbandbreite beseitigte, ließ es auf modernen GPUs noch einen erheblichen Teil der reinen Rechenleistung ungenutzt. FlashAttention-2 (2023) und FlashAttention-3 (2024) haben die Zusammenarbeit von GPU-Threads und Hardware-Einheiten neu gestaltet.

• FlashAttention-2 (bessere Arbeitsaufteilung).

• Weniger Nicht-Matmul-Operationen. GPU-Tensor-Cores führen Matrixmultiplikationen (matmul) etwa 16× schneller aus als herkömmliche CUDA-Kerne skalare Operationen wie die Softmax-Neuskalierung. FlashAttention-2 hat die Berechnung des Online-Softmax so umgeschrieben, dass die Neuskalierung erst ganz am Ende der inneren Schleife erfolgt. Dadurch fallen weniger langsame Nicht-Matmul-Anweisungen an.
• Parallelisierung über die Sequenzlänge. Die ursprüngliche Version parallelisierte die Arbeit hauptsächlich über Batchgröße und Attention-Heads. Beim Training mit langen Sequenzen und kleinen Batches blieben dadurch GPU-Kerne unbeschäftigt. FlashAttention-2 ergänzt Parallelismus über die Sequenzlängendimension und kehrt die Schleifenreihenfolge um (äußere Schleife über Q-Blöcke, innere Schleife über K und V), wodurch langsame Schreibzugriffe auf den gemeinsam genutzten Speicher minimiert werden.

• FlashAttention-3 (Nutzung der Hopper-H100-Hardware).

• Asynchrone Datenübertragung (TMA). NVIDIAs Hopper-Architektur führte den Tensor Memory Accelerator (TMA) ein, eine eigene Hardware-Einheit, die Daten zwischen globalem HBM und SRAM auf dem Chip kopiert, ohne Rechen-Threads zu belegen.
• Warp-Spezialisierung (Erzeuger-Verbraucher-Modell). FlashAttention-3 teilt GPU-Warps (Gruppen von 32 Threads) in feste Rollen auf. Erzeuger-Warps verwenden TMA, um künftige K- und V-Kacheln in den SRAM vorzuladen, während Verbraucher-Warps die aktuellen Kacheln an die Tensor Cores übergeben (WGMMA-Anweisungen). Dadurch überlagern sich Speicherübertragungen gleichzeitig mit Matrix- und Softmax-Berechnungen.
• Unterstützung geringer FP8-Präzision. FlashAttention-3 erschließt Hoppers FP8-Tensor-Cores und verdoppelt den Rohdurchsatz gegenüber FP16. Um Präzisionsverluste durch Ausreißer-Aktivierungen in FP8 zu verhindern, nutzt es Blockquantisierung (unabhängige Skalierung kleiner Kacheln) und inkohärente Verarbeitung (Multiplikation von Q und K mit zufälligen orthogonalen Matrizen, um Extremwerte vor der Quantisierung abzumildern).

2. Multi-GPU-Skalierungsstrategien (3D-Parallelismus)

Wenn ein Modell Dutzende oder Hunderte Milliarden Parameter hat, übersteigen seine Gewichte, Gradienten und Optimiererzustände den VRAM einer einzelnen GPU. Großskaliges Training kombiniert drei unterschiedliche Parallelisierungsstrategien, die oft als 3D-Parallelismus bezeichnet werden.

• Datenparallelismus & ZeRO / FSDP. Beim herkömmlichen Datenparallelismus hält jede GPU eine vollständige Kopie des Modells vor und verschwendet damit Speicher. Moderne Frameworks verwenden ZeRO (Zero Redundancy Optimizer) oder PyTorch FSDP (Fully Sharded Data Parallel). Dabei werden Parameter, Gradienten und Optimiererzustände auf alle GPUs aufgeteilt und die Gewichte jeder Schicht genau rechtzeitig für die Berechnung zusammengeführt.

• Tensorparallelismus (innerhalb einer Schicht). Statt eine riesige Feed-Forward-Matrixmultiplikation Y = XW zu duplizieren, teilt Tensorparallelismus W spaltenweise auf mehrere GPUs auf (W = [W_1, W_2]). Jede GPU berechnet gleichzeitig ihre Hälfte der Ausgabe und tauscht Teilergebnisse über schnelles NVLink aus (900 GB/s auf H100-GPUs).

• Pipelineparallelismus (zwischen Schichten). Die Schichten 1–8 liegen auf GPU 1, die Schichten 9–16 auf GPU 2 und so weiter. Damit nachgelagerte GPUs nicht untätig auf die Fertigstellung früherer Schichten warten (die sogenannte „Pipeline-Blase“), zerlegen Frameworks jeden Trainings-Batch in kleinere Micro-Batches, sodass alle GPUs wie an einem Fließband beschäftigt bleiben.

3. Quantisierung (GPTQ, AWQ) und Kernels zur Dequantisierung während der Ausführung

Während der Inferenz (der schrittweisen Erzeugung einzelner Tokens) werden große Sprachmodelle durch die Speicherbandbreite begrenzt, nicht durch die Rechenleistung. Das Laden von 70 Milliarden 16-Bit-Parametern (140 GB) aus dem VRAM für die Erzeugung eines einzigen Tokens lässt die Tensor Cores auf Speicherübertragungen warten. Quantisierung komprimiert Gewichte von 16-Bit-Gleitkommazahlen auf 4-Bit-Ganzzahlen (INT4) und verringert Speicherbedarf und Übertragungszeit um nahezu den Faktor 4×.

• GPTQ (Quantisierung nach dem Training durch Hesse-Kalibrierung).

• Naives Runden von 16-Bit-Gewichten auf 4-Bit-Ganzzahlen zerstört die Modellgenauigkeit. GPTQ führt einen kleinen Kalibrierungsdatensatz durch das Modell und misst mithilfe von Mathematik zweiter Ordnung (der inversen Hesse-Matrix), wie empfindlich jedes Gewicht ist.
• Die Gewichte werden spaltenweise quantisiert. Sobald das Runden eines Gewichts einen Fehler verursacht, passt GPTQ die übrigen, noch nicht quantisierten Gewichte in dieser Zeile mathematisch an, um den Fehler auszugleichen.

• AWQ (aktivierungsbewusste Gewichtsquantisierung).

• Forscher entdeckten, dass ungefähr 0,1 % bis 1 % der Gewichtskanäle weitaus wichtiger sind als die übrigen, weil sie große Aktivierungswerte multiplizieren.
• Statt diese entscheidenden Gewichte in 16 Bit zu belassen (was unübersichtliche, langsame Speicherlayouts mit gemischter Präzision erzeugt), skaliert AWQ diese besonders wichtigen Gewichtskanäle vor der Quantisierung hoch und die entsprechenden Eingangsaktivierungen zur Laufzeit herunter. So bleibt die Genauigkeit erhalten, während 100 % der Gewichte im einheitlichen 4-Bit-Format bleiben.

• Wie spezielle CUDA-Kernels während der Ausführung dequantisieren (W4A16).

• GPUs können eine 16-Bit-Aktivierung (FP16) nicht direkt mit einem ganzzahligen 4-Bit-Gewicht (INT4) multiplizieren. Stattdessen führen spezialisierte Inferenz-Kernels (etwa Marlin oder ExLlama) W4A16 aus (4-Bit-Gewichte, 16-Bit-Aktivierungen):
• 1. Komprimierter HBM-Lesezugriff: Der Kernel lädt gepackte 4-Bit-Gewichte (acht INT4-Werte in einer einzelnen 32-Bit-Ganzzahl) aus dem langsamen globalen VRAM und reduziert den Speicherbandbreitenbedarf um den Faktor 4×.
• 2. Entpacken auf Registerebene: In den ultraschnellen Registern der GPU verwendet der CUDA-Kernel schnelle bitweise Verschiebungen (>>) und Bitmasken (&), um die INT4-Ganzzahlen zu entpacken. Anschließend wendet er gespeicherte Skalierungs-/Nullpunktfaktoren an, um sie während der Ausführung wieder in FP16-Werte umzuwandeln.
• 3. Sofortige Tensor-Core-Berechnung: Die frisch entpackten FP16-Gewichte werden zusammen mit den FP16-Aktivierungen direkt in die Tensor Cores eingespeist, ohne die entpackten 16-Bit-Gewichte jemals zurück in den VRAM zu schreiben.