Transformador Distribuido de Secuencias Ultra-Largas

Ultra-Long Sequence Distributed Transformer

November 4, 2023

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

cs.AI

Resumen

Los modelos Transformer entrenados con secuencias largas suelen alcanzar una mayor precisión que aquellos entrenados con secuencias cortas. Sin embargo, los transformers convencionales enfrentan dificultades al entrenar secuencias largas debido a los abrumadores requisitos de computación y memoria. Los métodos existentes para el entrenamiento de secuencias largas ofrecen una aceleración y reducción de memoria limitadas, y pueden comprometer la precisión. Este artículo presenta un método novedoso y eficiente para el entrenamiento distribuido, el Transformer de Secuencia Larga-Corta (LSS Transformer), diseñado para entrenar transformers con secuencias largas. Este método distribuye una secuencia larga en segmentos entre las GPUs, donde cada GPU calcula una autoatención parcial para su segmento. Luego, utiliza una comunicación fusionada y una técnica novedosa de promediado doble de gradientes para evitar la necesidad de agregar las autoatenciones parciales y minimizar la sobrecarga de comunicación. Evaluamos el rendimiento del LSS Transformer en comparación con el paralelismo de secuencias de última generación de Nvidia utilizando el conjunto de datos Wikipedia enwik8. Los resultados muestran que nuestro método propuesto es 5.6 veces más rápido y 10.2 veces más eficiente en memoria en comparación con el paralelismo de secuencias de última generación en 144 GPUs Nvidia V100. Además, nuestro algoritmo escala a una longitud de secuencia extrema de 50,112 en 3,456 GPUs, logrando una eficiencia paralela superlineal del 161% y un rendimiento de 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.