Ultra-Long Sequence Distributed Transformer
Authors: name: Xiao Wang; name: Isaac Lyngaas; name: Aristeidis Tsaris; name: Peng Chen; name: Sajal Dash; name: Mayanka Chandra Shekar; name: Tao Luo; name: Hong-Jun Yoon; name: Mohamed Wahib; name: John Gouley
Published: 2023-11-04 11:38:53

Abstract: 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.