DeepSpeed Ulysses: Systeemoptimalisaties voor het mogelijk maken van training van Transformer-modellen met extreem lange sequenties

DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models

September 25, 2023

Auteurs: Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song, Samyam Rajbhandari, Yuxiong He

cs.AI

Samenvatting

Berekeningen in een typisch Transformer-gebaseerd groot taalmodel (LLM) kunnen worden gekarakteriseerd door batchgrootte, verborgen dimensie, aantal lagen en sequentielengte. Tot nu toe hebben systeemwerken voor het versnellen van LLM-training zich gericht op de eerste drie dimensies: dataparallelisme voor batchgrootte, tensorparallelisme voor verborgen grootte en pijplijnparallelisme voor modeldiepte of lagen. Deze breed bestudeerde vormen van parallelisme zijn niet gericht of geoptimaliseerd voor lange-sequentie Transformer-modellen. Gezien de praktische toepassingsbehoeften voor lange-sequentie LLM, wordt er opnieuw aandacht besteed aan sequentieparallelisme. Bestaande werken in sequentieparallelisme worden echter beperkt door geheugen-communicatie inefficiëntie, wat hun schaalbaarheid naar grote modellen met lange sequenties beperkt. In dit werk introduceren we DeepSpeed-Ulysses, een nieuwe, draagbare en effectieve methodologie voor het mogelijk maken van zeer efficiënte en schaalbare LLM-training met extreem lange sequentielengtes. DeepSpeed-Ulysses verdeelt in de kern invoergegevens langs de sequentiedimensie en gebruikt een efficiënte all-to-all collectieve communicatie voor aandachtberekening. Theoretische communicatieanalyse toont aan dat terwijl andere methoden communicatieoverhead veroorzaken naarmate de sequentielengte toeneemt, DeepSpeed-Ulysses een constant communicatievolume behoudt wanneer de sequentielengte en rekenapparaten proportioneel worden vergroot. Bovendien tonen experimentele evaluaties aan dat DeepSpeed-Ulysses 2,5 keer sneller traint met een 4 keer langere sequentielengte dan de bestaande SOTA-baseline.

English

Computation in a typical Transformer-based large language model (LLM) can be characterized by batch size, hidden dimension, number of layers, and sequence length. Until now, system works for accelerating LLM training have focused on the first three dimensions: data parallelism for batch size, tensor parallelism for hidden size and pipeline parallelism for model depth or layers. These widely studied forms of parallelism are not targeted or optimized for long sequence Transformer models. Given practical application needs for long sequence LLM, renewed attentions are being drawn to sequence parallelism. However, existing works in sequence parallelism are constrained by memory-communication inefficiency, limiting their scalability to long sequence large models. In this work, we introduce DeepSpeed-Ulysses, a novel, portable and effective methodology for enabling highly efficient and scalable LLM training with extremely long sequence length. DeepSpeed-Ulysses at its core partitions input data along the sequence dimension and employs an efficient all-to-all collective communication for attention computation. Theoretical communication analysis shows that whereas other methods incur communication overhead as sequence length increases, DeepSpeed-Ulysses maintains constant communication volume when sequence length and compute devices are increased proportionally. Furthermore, experimental evaluations show that DeepSpeed-Ulysses trains 2.5X faster with 4X longer sequence length than the existing method SOTA baseline.