FlashAttention-2: Schnellere Attention durch bessere Parallelisierung und Arbeitsaufteilung

• Sprachmodelle mit längerem Kontext wie GPT-4(32k), MPT(65k) und Claude(100k) sind erschienen
• Die Kontextlänge von Transformern zu erweitern ist schwierig, weil die Laufzeit- und Speicheranforderungen quadratisch ansteigen
• Das im vergangenen Jahr veröffentlichte FlashAttention wird vielerorts eingesetzt, da es den Speicherverbrauch senkt und die Attention-Geschwindigkeit erhöht
• Schon bei der Veröffentlichung war es 2- bis 4-mal schneller, aber es gab noch Raum für Verbesserungen. Im Vergleich zu optimierten Matrixmultiplikationen (GEMM) war es noch nicht schnell genug und erreichte theoretisch nur 25–40 % der maximalen FLOPs/s (maximal 124 TFLOPs/s auf einer A100 GPU)
• FlashAttention-2 ist doppelt so schnell wie die Vorgängerversion und liefert auf einer A100 GPU bis zu 230 TFLOPs/s Leistung
• Beim Training GPT-artiger Sprachmodelle wurden bis zu 225 TFLOPs erreicht (72 % Model-FLOP-Auslastung)
• Der Algorithmus wurde angepasst, um non-matmul FLOPs zu reduzieren
• Bessere Parallelisierung und eine geänderte Methode der Arbeitsaufteilung in jedem Thread-Block
• Die Anzahl der Head Dimensions wurde von 128 auf 256 erweitert