Ultra-Lang Sequentie Gedistribueerde Transformer

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

Samenvatting

Transformermodellen die getraind zijn op lange sequenties behalen vaak een hogere nauwkeurigheid dan korte sequenties. Helaas hebben conventionele transformers moeite met het trainen van lange sequenties vanwege de overweldigende reken- en geheugenvereisten. Bestaande methoden voor het trainen van lange sequenties bieden beperkte versnelling en geheugenreductie, en kunnen de nauwkeurigheid in gevaar brengen. Dit artikel presenteert een nieuwe en efficiënte gedistribueerde trainingsmethode, de Long Short-Sequence Transformer (LSS Transformer), voor het trainen van transformers met lange sequenties. Het verdeelt een lange sequentie in segmenten over GPU's, waarbij elke GPU een gedeeltelijke self-attention berekent voor zijn segment. Vervolgens gebruikt het een gefuseerde communicatie en een nieuwe dubbele gradient-averagingtechniek om de noodzaak om gedeeltelijke self-attention te aggregeren te vermijden en de communicatie-overhead te minimaliseren. We hebben de prestaties van de LSS Transformer vergeleken met de state-of-the-art Nvidia sequence-paralleliteit op een Wikipedia enwik8-dataset. De resultaten laten zien dat onze voorgestelde methode leidt tot een 5,6x snellere en 10,2x geheugenefficiëntere implementatie in vergelijking met state-of-the-art sequence-paralleliteit op 144 Nvidia V100 GPU's. Bovendien schaalt ons algoritme naar een extreme sequentielengte van 50.112 op 3.456 GPU's, waarbij het een super-lineaire parallelle efficiëntie van 161% en een doorvoer van 32 petaflops bereikt.

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.