Les Transformers à longues séquences étaient lents et gourmands en mémoire parce que l'auto-attention standard matérialisait de grandes matrices intermédiaires et les déplaçait à plusieurs reprises entre la mémoire à haute bande passante du GPU et les unités de calcul.
La solution tentante était l'attention approximative : modifier les mathématiques du modèle pour réduire la complexité. FlashAttention a pris la voie opposée. Il a conservé l'attention exacte, mais a traité des blocs dans la SRAM et recalculé les statistiques courantes du softmax afin que la matrice d'attention complète n'ait jamais à être stockée.
L'article a rapporté une accélération de 15 % du temps réel de bout en bout sur BERT-large par rapport au record d'entraînement MLPerf 1.1, une accélération de 3x sur GPT-2 à une longueur de séquence de 1K, et une accélération de 2,4x sur les tâches Long Range Arena.
Optimiser le trafic mémoire, pas les mathématiques
La contrainte centrale était les entrées-sorties entre la HBM et la SRAM, et pas seulement les opérations en virgule flottante.
Le découpage en tuiles permettait à chaque bloc d'effectuer un travail utile pendant que les données étaient déjà sur la puce.
Conserver l'algorithme exact évitait les régressions de qualité qui rendaient de nombreuses méthodes approximatives peu attrayantes.
La même primitive a ouvert des contextes plus longs et de nouvelles tâches de séquence sans modifier l'abstraction du Transformer.
Lorsque le déplacement sur le matériel est le goulot d'étranglement, un meilleur algorithme peut être celui qui calcule la même réponse en traversant moins souvent la hiérarchie mémoire.
Suites
FlashAttention est devenu un bloc de construction standard pour l'entraînement et l'inférence efficaces des Transformers. Sa leçon plus large s'est répandue au-delà de l'attention : les algorithmes modernes pour accélérateurs gagnent souvent en concevant autour de la hiérarchie mémoire avant d'inventer de nouvelles mathématiques de modèle.
VÉRIFIEZ LES SOURCES
Les sources
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness arXiv · arxiv.org ·