Трансформер с распределенной обработкой сверхдлинных последовательностей

Ultra-Long Sequence Distributed Transformer

November 4, 2023

Авторы: Xiao Wang, Isaac Lyngaas, Aristeidis Tsaris, Peng Chen, Sajal Dash, Mayanka Chandra Shekar, Tao Luo, Hong-Jun Yoon, Mohamed Wahib, John Gouley

cs.AI

Аннотация

Трансформерные модели, обученные на длинных последовательностях, часто достигают более высокой точности, чем модели, обученные на коротких последовательностях. К сожалению, традиционные трансформеры сталкиваются с трудностями при обучении на длинных последовательностях из-за чрезмерных требований к вычислениям и памяти. Существующие методы для обучения на длинных последовательностях предлагают ограниченное ускорение и снижение потребления памяти, а также могут снижать точность. В данной статье представлен новый и эффективный метод распределенного обучения — Long Short-Sequence Transformer (LSS Transformer), предназначенный для обучения трансформеров на длинных последовательностях. Он распределяет длинную последовательность на сегменты между GPU, где каждый GPU вычисляет частичное самовнимание для своего сегмента. Затем используется объединенная коммуникация и новая техника двойного усреднения градиентов, чтобы избежать необходимости агрегирования частичного самовнимания и минимизировать накладные расходы на коммуникацию. Мы оценили производительность LSS Transformer в сравнении с современным методом параллелизма последовательностей от Nvidia на наборе данных Wikipedia enwik8. Результаты показывают, что предложенный метод обеспечивает ускорение в 5,6 раз и снижение потребления памяти в 10,2 раза по сравнению с современным методом параллелизма последовательностей на 144 GPU Nvidia V100. Более того, наш алгоритм масштабируется до экстремальной длины последовательности в 50 112 на 3 456 GPU, достигая 161% сверхлинейной параллельной эффективности и пропускной способности в 32 петафлопса.

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.