Trasformatore Distribuito per Sequenze Ultra-Lunghe

Ultra-Long Sequence Distributed Transformer

November 4, 2023

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

cs.AI

Abstract

I modelli Transformer addestrati su sequenze lunghe spesso raggiungono una maggiore accuratezza rispetto a sequenze brevi. Sfortunatamente, i transformer convenzionali incontrano difficoltà nell'addestramento su sequenze lunghe a causa degli eccessivi requisiti di calcolo e memoria. I metodi esistenti per l'addestramento su sequenze lunghe offrono un limitato aumento di velocità e riduzione della memoria, e possono compromettere l'accuratezza. Questo articolo presenta un metodo innovativo ed efficiente per l'addestramento distribuito, il Long Short-Sequence Transformer (LSS Transformer), progettato per addestrare transformer su sequenze lunghe. Esso suddivide una sequenza lunga in segmenti distribuiti tra le GPU, con ciascuna GPU che calcola un'attenzione parziale (self-attention) per il proprio segmento. Successivamente, utilizza una comunicazione fusa e una nuova tecnica di mediazione doppia del gradiente per evitare la necessità di aggregare le attenzioni parziali e minimizzare l'overhead di comunicazione. Abbiamo valutato le prestazioni tra LSS Transformer e il parallelismo di sequenza all'avanguardia di Nvidia su un dataset Wikipedia enwik8. I risultati mostrano che il nostro metodo proposto porta a un'implementazione 5,6 volte più veloce e 10,2 volte più efficiente in termini di memoria rispetto al parallelismo di sequenza all'avanguardia su 144 GPU Nvidia V100. Inoltre, il nostro algoritmo scala fino a una lunghezza estrema della sequenza di 50.112 su 3.456 GPU, raggiungendo un'efficienza parallela super-lineare del 161% e una velocità di elaborazione di 32 petaflop.

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.