DeepSpeed Ulysses: Ottimizzazioni di sistema per abilitare l'addestramento di modelli Transformer con sequenze estremamente lunghe

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

September 25, 2023

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

cs.AI

Abstract

Il calcolo in un tipico modello linguistico di grandi dimensioni (LLM) basato su Transformer può essere caratterizzato da dimensione del batch, dimensione nascosta, numero di livelli e lunghezza della sequenza. Fino ad ora, i lavori di sistema per accelerare l'addestramento degli LLM si sono concentrati sui primi tre aspetti: parallelismo dei dati per la dimensione del batch, parallelismo tensoriale per la dimensione nascosta e parallelismo pipeline per la profondità del modello o i livelli. Queste forme di parallelismo ampiamente studiate non sono mirate o ottimizzate per modelli Transformer con sequenze lunghe. Date le esigenze pratiche delle applicazioni per LLM con sequenze lunghe, si sta ridestando l'attenzione sul parallelismo delle sequenze. Tuttavia, i lavori esistenti sul parallelismo delle sequenze sono limitati da inefficienze di memoria-comunicazione, che ne riducono la scalabilità per modelli di grandi dimensioni con sequenze lunghe. In questo lavoro, introduciamo DeepSpeed-Ulysses, una metodologia innovativa, portabile ed efficace per abilitare un addestramento di LLM altamente efficiente e scalabile con lunghezze di sequenza estremamente elevate. DeepSpeed-Ulysses, nel suo nucleo, partiziona i dati di input lungo la dimensione della sequenza e utilizza una comunicazione collettiva all-to-all efficiente per il calcolo dell'attenzione. L'analisi teorica della comunicazione mostra che, mentre altri metodi comportano un sovraccarico di comunicazione all'aumentare della lunghezza della sequenza, DeepSpeed-Ulysses mantiene un volume di comunicazione costante quando la lunghezza della sequenza e i dispositivi di calcolo vengono aumentati proporzionalmente. Inoltre, le valutazioni sperimentali dimostrano che DeepSpeed-Ulysses addestra 2,5 volte più velocemente con una lunghezza di sequenza 4 volte maggiore rispetto al metodo esistente di riferimento SOTA.

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.