Qu'est-ce que FlashAttention et quel gain de mémoire offre-t-il?

FlashAttention est un algorithme optimisé pour le calcul du mécanisme d'attention dans les transformateurs, qui réduit considérablement la consommation de mémoire et accélère les calculs par rapport à l'approche classique.

L'idée principale de FlashAttention est de calculer l'attention sans stocker explicitement de grandes matrices de résultats intermédiaires (par exemple, la matrice de poids d'attention QK^T), en utilisant un traitement par blocs et une utilisation efficace de la mémoire cache GPU. Cela permet de réduire la consommation maximale de mémoire d'environ 3-4 fois, ce qui est particulièrement important lors de l'entraînement de grands modèles de langage (LLM) avec de longues séquences.

Environ, si l'attention classique nécessite une mémoire O(N^2) pour une séquence de longueur N, FlashAttention réduit cette charge grâce à des calculs à la volée et des optimisations, permettant de faire fonctionner des modèles avec des contextes plus grands sur les mêmes GPU.

Ainsi, FlashAttention est une méthode pour rendre l'attention plus efficace en mémoire et en vitesse, ce qui est critique pour les modèles à grande échelle.