Ces dernières années, les grands modèles de langage (LLM) basés sur l'architecture Transformer ont suscité une large attention grâce à leurs performances remarquables. Les LLM industriels doivent traiter des séquences longues pour fournir des services de haute qualité. Toutefois, la consommation de mémoire augmente de manière quadratique avec la longueur de la séquence, ce qui limite la capacité d'extension de l'entraînement sur des séquences longues. Les méthodes parallèles existantes génèrent des tenseurs redondants pendant l'exécution, offrant ainsi un potentiel d'optimisation de la mémoire ; par ailleurs, le parallélisme par tenseur (TP) ne permet pas un chevauchement efficace entre le calcul et la communication. Pour résoudre ces problèmes, cet article propose une méthode parallèle universelle — le parallélisme de tenseur à mémoire efficace (METP), conçue spécifiquement pour les unités de calcul clés dans l'entraînement des Transformers (à savoir deux multiplications matricielles consécutives et la fonction potentielle entre elles O=f(AB)C). METP répartit les sous-tâches de calcul de O sur plusieurs dispositifs, en utilisant une communication point-à-point (send/recv) au lieu d'une communication collective pour échanger les sous-matrices nécessaires, évitant ainsi la création de tenseurs redondants. Grâce à la technique de double buffering, un chevauchement profond entre calcul et communication est réalisé, tout en proposant des conditions théoriques de chevauchement complet pour guider l'entraînement des Transformers sur de longues séquences. L'analyse théorique montre que pour un parallélisme de degré p, le coût mémoire de METP sans FlashAttention pour le calcul de l'attention est O(1/p3); lors de l'utilisation de FlashAttention pour le calcul de l'attention multi-tête, METP peut économiser au moins 41,7 % de mémoire par rapport au TP. Les expériences démontrent que, avec une configuration de 8 GPU A100, METP peut augmenter la longueur des séquences de 2,38 à 2,99 fois par rapport à d'autres méthodes.
Keywords
apprentissage distribué;grands modèles de langage;séquences longues;systèmes d'apprentissage automatique;efficacité mémoire;parallélisme de tenseur