Parallélisme des séquences d'Ulysse : formation avec des contextes d'un million de jetons

Entraîner de grands modèles de langage sur de longues séquences est devenu essentiel pour créer des systèmes d’IA performants. Alors que les modèles sont de plus en plus utilisés pour des tâches telles que l’analyse de documents, la compréhension du code, le raisonnement complexe et les charges de travail RAG, la nécessité de traiter des séquences de centaines de milliers, voire de millions, de jetons s’est considérablement accrue. Pour mettre cela en perspective, un livre moyen contient environ 250 000 jetons, donc la formation sur des contextes multi-documents ou des entrées de longueur de livre nécessite de gérer des séquences bien au-delà de ce qui peut tenir sur un seul GPU. Cependant, l'entraînement avec des contextes aussi longs présente des défis de mémoire importants : le calcul de l'attention évolue quadratiquement avec la longueur de la séquence, dépassant rapidement la mémoire du GPU pour des contextes dépassant des dizaines de milliers de jetons.

Le parallélisme de séquence d'Ulysses (qui fait partie du protocole Arctic Long Sequence Training (ALST) de Snowflake AI Research) fournit une solution élégante en répartissant le calcul de l'attention sur plusieurs GPU via le parallélisme des têtes d'attention. Dans cet article, nous explorerons le fonctionnement d'Ulysses et la manière dont il a été intégré dans l'écosystème Hugging Face, de Accelerate au Transformers Trainer en passant par le SFTTrainer de TRL.

Contenu

• Le défi de l’entraînement en séquences longues
• Comment fonctionne Ulysse
• Intégration avec Accelerate
• Intégration avec Transformers Trainer
• Intégration avec SFTTrainer de TRL
• Comparaison d'Ulysse et de Ring Attention
• Meilleures pratiques
• Repères
• Ressources

Le défi de l’entraînement en séquences longues

Le mécanisme d'attention dans les transformateurs évolue quadratiquement avec la longueur de la séquence. Pour une séquence de longueur n n n, l'attention standard nécessite une mémoire O(n2) O(n^2) O(n2) FLOP et O(n2) O(n^2) O(n2) pour calculer et stocker la matrice de score d'attention. Les implémentations optimisées comme FlashAttention réduisent la mémoire à O(n) O(n) O(n) en mosaïque le calcul et en ne matérialisant jamais la matrice d'attention complète, mais le calcul O(n2) O(n^2) O(n2) reste. Pour les séquences très longues (plus de 32 000 jetons), même avec FlashAttention, la formation repousse toujours les limites de la mémoire mono-GPU.

Considérez ces scénarios dans lesquels une formation en contexte long est essentielle :

• Compréhension de documents : traitement de livres entiers, de documents juridiques ou de documents de recherche
• Analyse de code : Comprendre les bases de code volumineuses avec plusieurs fichiers interconnectés
• Tâches de raisonnement : les modèles qui « réfléchissent » étape par étape peuvent générer des milliers de jetons lors de l'inférence
• Génération augmentée par récupération : intégration de nombreux passages récupérés dans le contexte

Le parallélisme traditionnel des données n'aide pas ici : chaque GPU doit toujours traiter la séquence complète à l'intérieur du bloc d'attention. Nous avons besoin d'un moyen de diviser la séquence elle-même sur plusieurs appareils.

Comment fonctionne Ulysse

Le parallélisme de séquence d'Ulysses (SP), introduit dans l'article DeepSpeed ​​Ulysses, adopte une approche intelligente : en plus de diviser la dimension de la séquence, il divise également l'attention entre les GPU.

Ulysses divise les séquences d'entrée selon la dimension de la séquence et utilise la communication tout-à-tout pour échanger des paires clé-valeur, permettant à chaque GPU de calculer un sous-ensemble de têtes d'attention. (Source : Blog d'ingénierie des flocons de neige)

Voici comment cela fonctionne :

1. Partage de séquence : la séquence d'entrée est divisée le long de la dimension de séquence sur les GPU P P P. Chaque GPU i i i contient des jetons [i⋅n/P,(i+1)⋅n/P) [i \cdot n/P, (i+1) \cdot n/P) [i⋅n/P,(i+1)⋅n/P).

2. Projection QKV : chaque GPU calcule les projections de requête, de clé et de valeur pour son morceau de séquence locale.

3. Communication tout-à-tout : une opération collective tout-à-tout redistribue les données de sorte que chaque GPU conserve toutes les positions de séquence après les proj...
   [Courte citation de 8% de l'article original]