Ultra-Lange Sequenzverteilte Transformer

Ultra-Long Sequence Distributed Transformer

November 4, 2023

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

cs.AI

Zusammenfassung

Transformer-Modelle, die mit langen Sequenzen trainiert werden, erreichen oft eine höhere Genauigkeit als solche mit kurzen Sequenzen. Leider haben konventionelle Transformer Schwierigkeiten beim Training mit langen Sequenzen aufgrund der überwältigenden Rechen- und Speicheranforderungen. Bestehende Methoden für das Training mit langen Sequenzen bieten nur begrenzte Beschleunigung und Speicherreduzierung und können die Genauigkeit beeinträchtigen. Dieses Papier stellt eine neuartige und effiziente Methode für das verteilte Training vor, den Long Short-Sequence Transformer (LSS Transformer), der das Training von Transformern mit langen Sequenzen ermöglicht. Dabei wird eine lange Sequenz in Segmente aufgeteilt, die auf GPUs verteilt werden, wobei jede GPU eine partielle Selbstaufmerksamkeit für ihr Segment berechnet. Anschließend wird eine kombinierte Kommunikation und eine neuartige Technik zur doppelten Gradientenmittelung verwendet, um die Notwendigkeit der Aggregation partieller Selbstaufmerksamkeit zu vermeiden und den Kommunikationsaufwand zu minimieren. Wir haben die Leistung des LSS Transformers im Vergleich zum state-of-the-art Nvidia Sequence Parallelism auf dem Wikipedia enwik8-Datensatz evaluiert. Die Ergebnisse zeigen, dass unsere vorgeschlagene Methode im Vergleich zum state-of-the-art Sequence Parallelism auf 144 Nvidia V100 GPUs eine 5,6-fache Beschleunigung und eine 10,2-fache Speichereffizienz erreicht. Darüber hinaus skaliert unser Algorithmus auf eine extreme Sequenzlänge von 50.112 bei 3.456 GPUs und erreicht dabei eine superlineare parallele Effizienz von 161 % und einen Durchsatz von 32 Petaflops.

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.