FlashAttention-3 optimise les GPU H100 pour améliorer les performances des grands modèles de langage.

L'attention est un élément fondamental de l'architecture transformer qui propulse les grands modèles de langage (LLMs). Cependant, à mesure que les LLMs s'étendent et gèrent des séquences d'entrée de plus en plus longues, le coût computationnel de l'attention devient un goulot d'étranglement significatif. Pour résoudre ce problème, une équipe collaborative de Colfax Research, Meta, Nvidia, Georgia Tech, Princeton University et Together AI a introduit FlashAttention-3. Cette technique de pointe accélére substantiellement le calcul de l'attention sur les GPU Nvidia Hopper (H100 et H800).

Comprendre les Défis du Calcul de l'Attention dans les LLMs

Le mécanisme d'attention du modèle transformer permet d'évaluer les relations entre différents tokens dans une séquence d'entrée. Bien que efficace, ce mécanisme est computationnellement intensif. À mesure que la longueur des séquences d'entrée augmente, le coût du calcul de l'attention croît de manière quadratique, entraînant des goulots d'étranglement significatifs pour l'évolutivité des LLMs. De plus, les GPU modernes sont principalement optimisés pour les opérations de multiplication matricielle, qui sont au cœur du deep learning. D'autres opérations, comme l'exponentiation, sont beaucoup plus lentes, aggravant les problèmes de performance.

Amélioration de l'Utilisation des Ressources Matérielles avec FlashAttention

FlashAttention, lancé en 2022, a amélioré l'efficacité des calculs d'attention en minimisant les transferts de mémoire entre la mémoire à large bande (HBM) et la mémoire statique (SRAM) sur les GPU. En traitant les poids d'attention par petits morceaux ou "tuiles", FlashAttention a permis aux LLMs d'élargir leurs fenêtres de contexte de milliers à potentiellement millions de tokens. Cependant, avec l'avancement des capacités matérielles, le besoin d'optimisations supplémentaires s'est également imposé. FlashAttention-2, présenté en 2023, a optimisé l'utilisation des ressources GPU, atteignant 70 % de la performance maximale sur les GPU Nvidia A100, mais n'exploitant que 35 % des capacités de l'H100.

Innovations de FlashAttention-3

FlashAttention-3 tire parti des nouvelles fonctionnalités des GPU Nvidia Hopper pour améliorer les performances, notamment un meilleur débit pour les multiplications matricielles et des transferts de données plus rapides entre les segments de mémoire. Cela permet une efficacité accrue avec des opérations à faible précision. Parmi les innovations clés de FlashAttention-3, on trouve :

1. Planification Optimisée : Les opérations sont organisées pour maximiser le chevauchement entre le calcul et le transfert de données, réduisant le temps d'inactivité des GPU.

2. Interopérabilité Transparent : En combinant les opérations de multiplication matricielle et de softmax, FlashAttention-3 minimise les goulots d'étranglement potentiels.

3. Amélioration des Performances des Modèles Quantifiés : Des arrangements spéciaux dans les opérations garantissent des calculs plus rapides et plus précis, même avec des représentations à faible bit pour réduire la taille des modèles, tout en répondant à la problématique de l'exactitude généralement associée à la quantification.

Des recherches indiquent que FlashAttention-3 peut utiliser jusqu'à 75 % de la performance maximale du GPU H100, offrant un gain de vitesse de 1,5 à 2 fois par rapport aux versions précédentes de FlashAttention pour l'entraînement et le déploiement des LLMs.

Avantages de FlashAttention-3

L'accélération des calculs d'attention permise par FlashAttention-3 a d'importantes implications pour le développement et les applications des LLM :

- Entraînement Accéléré : L'efficacité accrue peut réduire considérablement les temps d'entraînement, permettant aux chercheurs d'explorer des modèles et des ensembles de données plus importants.

- Fenêtres de Contexte Élargies : En facilitant le traitement de séquences plus longues, FlashAttention-3 ouvre de nouvelles avenues pour les applications des LLM, telles que la compréhension de documents longs et l'apprentissage contextuel à nombreux exemples.

- Efficacité des Coûts : Une utilisation accrue des GPU peut réduire le nombre d'accélérateurs nécessaires pour les opérations LLM, diminuant ainsi les coûts de production.

FlashAttention-3 a été open-sourcé sous une licence permissive, avec des projets d'intégration dans des bibliothèques de deep learning populaires comme PyTorch et Hugging Face Transformers, permettant aux chercheurs et développeurs de tirer parti des avancées de FlashAttention-3. Comme le soulignent les chercheurs dans un article de blog de Together AI, « Concevoir des algorithmes qui exploitent les caractéristiques matérielles peut engendrer des améliorations significatives d'efficacité et libérer de nouvelles capacités pour les modèles. » Ils anticipent d'autres optimisations pour l'inférence des LLM et l'application de leurs techniques sur diverses architectures matérielles.