FlashAttention-2 : une attention plus rapide grâce à une meilleure parallélisation et à une meilleure répartition des tâches

• Des modèles de langage avec un contexte plus long ont fait leur apparition, comme GPT-4 (32k), MPT (65k) et Claude (100k)
• Étendre la longueur de contexte des transformers est difficile, car les besoins en temps d’exécution et en mémoire augmentent au carré
• FlashAttention, lancé l’an dernier, a été utilisé dans de nombreux contextes en réduisant l’usage mémoire et en augmentant la vitesse de l’attention
• Au moment de sa sortie, il était déjà 2 à 4 fois plus rapide, mais il restait encore une marge d’amélioration. Par rapport à une multiplication de matrices optimisée (GEMM), il n’était toujours pas aussi rapide et n’atteignait que 25 à 40 % des FLOPs/s théoriques maximums (jusqu’à 124 TFLOPs/s sur GPU A100)
• FlashAttention-2 est 2 fois plus rapide que la version précédente et offre jusqu’à 230 TFLOP/s sur GPU A100
• Lors de l’entraînement de modèles de langage de type GPT, il a atteint jusqu’à 225 TFLOPS (72 % d’utilisation des FLOP du modèle)
• L’algorithme a été ajusté afin de réduire les FLOPs hors multiplication de matrices
• Meilleure parallélisation et modification de la méthode de répartition du travail dans chaque bloc de threads
• Le nombre de Head Dimensions a été étendu de 128 à 256