[論文レビュー] DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models

Sam Adé Jacobs, Masahiro Tanaka+5 more|arXiv (Cornell University)|Sep 25, 2023

Topic Modeling被引用数 5

ひとこと要約

DeepSpeed-Ulyssesは、入力シーケンスをGPU間で分割し、効率的なall-to-all通信を用いることで、Transformerモデルの極めて長いシーケンス長（最大100万トークン）での学習を可能にする、革新的なシーケンス並列処理手法を導入している。従来のシステムと比較して2.5倍の高速化と10倍以上の通信削減を達成しながら、高いスループットとFlashAttentionおよびZeRO-3との互換性を維持している。

ABSTRACT

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.

研究の動機と目的

• Transformerモデルにおける極めて長いシーケンス長の処理に対応するための、従来の並列処理技術（データ並列、テンソル並列、パイプライン並列）のスケーラビリティ制限を解消すること。
• 長期間にわたりスケーラブルでないメモリ-通信非効率性を示す、従来のシーケンス並列処理アプローチの課題を克服すること。
• 既存のトレーニングフレームワークへの最小限のコード変更で実現可能な、汎用的でポータブルかつ最小限のコード変更で実現可能なシステム最適化を設計し、多様なハードウェアおよびアテンションメカニズムにおける長期間シーケンスLLMの効率的かつスケーラブルな学習を可能にすること。
• ピーク性能の54％を超える計算効率（計算利用率）を達成し、密度型およびスパース型アテンションの両方をサポートしながら、性能劣化を最小限に抑えること。

提案手法

• 複数のGPUにわたってシーケンス次元に沿って入力シーケンスを分割し、重複のないシーケンスセグメントを並列処理可能にする。
• アテンション計算の前後で、クエリ、キー、バリューをアテンションヘッド全体にわたって再配分するためのall-to-all通信プリミティブを採用し、各GPUが固有のヘッドサブセットに対して計算を行うことを保証する。
• 2段階のall-to-all通信を実施：第1段階では、各GPUが割り当てられたアテンションヘッド用の完全なシーケンスを受信できるようにデータを再構成し、第2段階では結果を集約しながらシーケンス次元に沿って再分割する。
• モデル並列処理およびメモリ最適化のためのZeRO-3とシームレスに統合され、大規模なモデルサイズと長期間シーケンス長の両方をサポートする。
• FlashAttention v2 やスパースアテンション（例：ブロックスパース）を含む効率的なアテンションメカニズムをサポートし、通信オーバーヘッドを低く保つ。
• 実装に依存しない設計を採用し、既存のトレーニングフレームワークへの最小限のコード変更で実現可能である。

Figure 1: Multi-head attention Transformer

実験結果

リサーチクエスチョン

• RQ1 特にシーケンス長が延びる際、シーケンス並列処理がスケーリングに耐えうる通信効率性を達成できるか。
• RQ2 シーケンス長とデバイス数が比例して増加する中で、all-to-all通信を活用することで通信量を一定に保てるか。
• RQ3 DeepSpeed-Ulyssesのようなシステムレベルの最適化が、侵襲的なコード再構築を要せず、長期間シーケンス学習における高いスループットと強力なスケーリングを達成できるか。
• RQ4 従来のシーケンス並列処理手法（例：Megatron-LM）と比較して、この手法がスループット、メモリ効率、スケーラビリティの観点でどの程度優れているか。
• RQ5 ZeRO-3およびFlashAttentionとの統合が、トレーニング効率とモデルサイズのスケーラビリティにどのような影響を与えるか。

主な発見

• DeepSpeed-Ulyssesは、最大100万トークンのシーケンス長のモデル学習を可能にし、既存のシステムと比較して4倍長いシーケンス長を処理できる。
• 密度型アテンションを用いた7Bおよび30Bパラメータモデルにおいて、SOTAベースライン（Megatron-LM）と比較して2.5倍の高いトレーニングスループットを達成している。
• シーケンス長とGPU数が比例して増加する中でも、通信量が一定に保たれる。これに対して、他の手法は通信オーバーヘッドの増加を伴う。
• 実験全体で1GPUあたり175 TFlopsを超える性能を維持し、ハードウェアピーク性能の54％以上を達成している。
• ZeRO-3統合によるメモリ節約のおかげで、同じハードウェア環境下でもMegatron-LMと比較して4倍の長期間シーケンス長が可能になっている。
• 強力な弱スケーリングを示しており、シーケンス長とGPU数が比例して増加してもスループットが高く維持される。通信オーバーヘッドによる性能劣化はわずかである。

Figure 2: DeepSpeed sequence parallelism (DeepSpeed-Ulysses) design