Transformador Distribuído para Sequências Ultra-Longas

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

Resumo

Modelos Transformer treinados em sequências longas frequentemente alcançam maior precisão do que em sequências curtas. Infelizmente, os transformers convencionais enfrentam dificuldades no treinamento de sequências longas devido aos requisitos computacionais e de memória excessivos. Os métodos existentes para treinamento de sequências longas oferecem aceleração e redução de memória limitadas, podendo comprometer a precisão. Este artigo apresenta um método de treinamento distribuído novo e eficiente, o Long Short-Sequence Transformer (LSS Transformer), para treinar transformers com sequências longas. Ele distribui uma sequência longa em segmentos entre GPUs, com cada GPU calculando uma autoatenção parcial para seu segmento. Em seguida, utiliza uma comunicação fundida e uma nova técnica de média dupla de gradientes para evitar a necessidade de agregar a autoatenção parcial e minimizar a sobrecarga de comunicação. Avaliamos o desempenho entre o LSS Transformer e o paralelismo de sequência de última geração da Nvidia em um conjunto de dados Wikipedia enwik8. Os resultados mostram que nosso método proposto leva a uma implementação 5,6 vezes mais rápida e 10,2 vezes mais eficiente em memória em comparação com o paralelismo de sequência de última geração em 144 GPUs Nvidia V100. Além disso, nosso algoritmo escala para um comprimento de sequência extremo de 50.112 em 3.456 GPUs, alcançando 161% de eficiência paralela superlinear e uma taxa de transferência 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.