Transformeur Distribué pour Séquences Ultra-Longues

Ultra-Long Sequence Distributed Transformer

November 4, 2023

Auteurs: Xiao Wang, Isaac Lyngaas, Aristeidis Tsaris, Peng Chen, Sajal Dash, Mayanka Chandra Shekar, Tao Luo, Hong-Jun Yoon, Mohamed Wahib, John Gouley

cs.AI

Résumé

Les modèles Transformer entraînés sur de longues séquences atteignent souvent une précision supérieure à celle des courtes séquences. Malheureusement, les Transformers conventionnels rencontrent des difficultés avec l'entraînement sur de longues séquences en raison des exigences écrasantes en calcul et en mémoire. Les méthodes existantes pour l'entraînement sur de longues séquences offrent une accélération et une réduction de mémoire limitées, et peuvent compromettre la précision. Cet article présente une nouvelle méthode d'entraînement distribué efficace, le Transformer à Longues et Courtes Séquences (LSS Transformer), pour l'entraînement de Transformers sur de longues séquences. Il répartit une longue séquence en segments parmi les GPU, chaque GPU calculant une auto-attention partielle pour son segment. Ensuite, il utilise une communication fusionnée et une nouvelle technique de moyennage double des gradients pour éviter la nécessité d'agréger les auto-attentions partielles et minimiser la surcharge de communication. Nous avons évalué les performances entre le LSS Transformer et le parallélisme de séquence de pointe de Nvidia sur un ensemble de données Wikipedia enwik8. Les résultats montrent que notre méthode proposée conduit à une implémentation 5,6 fois plus rapide et 10,2 fois plus économe en mémoire par rapport au parallélisme de séquence de pointe sur 144 GPU Nvidia V100. De plus, notre algorithme s'adapte à une longueur de séquence extrême de 50 112 sur 3 456 GPU, atteignant une efficacité parallèle super-linéaire de 161 % et un débit de 32 pétaflops.

English

Transformer models trained on long sequences often achieve higher accuracy than short sequences. Unfortunately, conventional transformers struggle with long sequence training due to the overwhelming computation and memory requirements. Existing methods for long sequence training offer limited speedup and memory reduction, and may compromise accuracy. This paper presents a novel and efficient distributed training method, the Long Short-Sequence Transformer (LSS Transformer), for training transformer with long sequences. It distributes a long sequence into segments among GPUs, with each GPU computing a partial self-attention for its segment. Then, it uses a fused communication and a novel double gradient averaging technique to avoid the need to aggregate partial self-attention and minimize communication overhead. We evaluated the performance between LSS Transformer and the state-of-the-art Nvidia sequence parallelism on a Wikipedia enwik8 dataset. Results show that our proposed method lead to 5.6x faster and 10.2x more memory-efficient implementation compared to state-of-the-art sequence parallelism on 144 Nvidia V100 GPUs. Moreover, our algorithm scales to an extreme sequence length of 50,112 at 3,456 GPUs, achieving 161% super-linear parallel efficiency and a throughput of 32 petaflops.