Langsequenz-Transformer waren langsam und speicherhungrig, weil die standardmäßige Self-Attention große Zwischenmatrizen materialisierte und sie wiederholt zwischen dem High-Bandwidth-Speicher der GPU und den Recheneinheiten hin- und herbewegte.
Der naheliegende Ausweg war approximative Attention: die Mathematik des Modells zu ändern, um die Komplexität zu senken. FlashAttention ging den entgegengesetzten Weg. Es behielt exakte Attention bei, verarbeitete aber Blöcke im SRAM und berechnete die laufenden Softmax-Statistiken neu, sodass die vollständige Attention-Matrix nie gespeichert werden musste.
Die Arbeit berichtete eine Ende-zu-Ende-Wall-Clock-Beschleunigung von 15 % bei BERT-large gegenüber dem MLPerf-1.1-Trainingsrekord, eine 3-fache Beschleunigung bei GPT-2 mit Sequenzlänge 1K und eine 2,4-fache Beschleunigung bei Long Range Arena-Aufgaben.
Speicherverkehr optimieren, nicht die Mathematik
Die zentrale Einschränkung war der IO zwischen HBM und SRAM, nicht nur die Gleitkommaoperationen.
Durch Kachelung konnte jeder Block nützliche Arbeit verrichten, während die Daten bereits auf dem Chip lagen.
Die Beibehaltung des exakten Algorithmus vermied die Qualitätseinbußen, die viele approximative Methoden unattraktiv machten.
Dieselbe Grundprimitive ermöglichte längere Kontexte und neue Sequenzaufgaben, ohne die Transformer-Abstraktion zu verändern.
Wenn die Hardware-Bewegung der Engpass ist, kann ein besserer Algorithmus derjenige sein, der dieselbe Antwort berechnet und dabei die Speicherhierarchie seltener durchläuft.
FlashAttention wurde zu einem Standardbaustein für effizientes Transformer-Training und -Inferenz. Seine größere Lehre reichte über Attention hinaus: Moderne Beschleuniger-Algorithmen gewinnen oft, indem sie um die Speicherhierarchie herum entwerfen, bevor sie neue Modellmathematik erfinden.
DEN BELEGEN FOLGEN
Die Quellen
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness arXiv · arxiv.org ·