초장거리 시퀀스 분산 트랜스포머

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가 해당 세그먼트에 대한 부분적 self-attention을 계산합니다. 그런 다음, 융합된 통신과 새로운 이중 그래디언트 평균화 기술을 사용하여 부분적 self-attention을 집계할 필요를 없애고 통신 오버헤드를 최소화합니다. 우리는 LSS Transformer와 최신 Nvidia 시퀀스 병렬화의 성능을 Wikipedia enwik8 데이터셋에서 평가했습니다. 결과는 제안된 방법이 144개의 Nvidia V100 GPU에서 최신 시퀀스 병렬화보다 5.6배 빠르고 10.2배 더 메모리 효율적인 구현을 이끌어냄을 보여줍니다. 또한, 우리의 알고리즘은 3,456개의 GPU에서 50,112의 극단적인 시퀀스 길이로 확장 가능하며, 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.