[論文レビュー] TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models

Zhuohan Li, Siyuan Zhuang+5 more|arXiv (Cornell University)|Feb 16, 2021

Topic Modeling参考文献 31被引用数 33

ひとこと要約

TeraPipeはTransformer言語モデルに対してトークンレベルのパイプライン並列性を導入し、AWS上の従来の同期モデル並列法に比べてGPT-3-175Bで最大5.0倍の訓練速度向上を達成する。

ABSTRACT

Model parallelism has become a necessity for training modern large-scale deep language models. In this work, we identify a new and orthogonal dimension from existing model parallel approaches: it is possible to perform pipeline parallelism within a single training sequence for Transformer-based language models thanks to its autoregressive property. This enables a more fine-grained pipeline compared with previous work. With this key idea, we design TeraPipe, a high-performance token-level pipeline parallel algorithm for synchronous model-parallel training of Transformer-based language models. We develop a novel dynamic programming-based algorithm to calculate the optimal pipelining execution scheme given a specific model and cluster configuration. We show that TeraPipe can speed up the training by 5.0x for the largest GPT-3 model with 175 billion parameters on an AWS cluster with 48 p3.16xlarge instances compared with state-of-the-art model-parallel methods. The code for reproduction can be found at https://github.com/zhuohan123/terapipe

研究の動機と目的

• 単一デバイスのメモリ制限を超えて極めて大規模なTransformer LMを訓練するために、より深いモデル並列性の必要性を喚起する。
• 自己回帰の依存関係を活用するトークン列に沿った新しい、細粒度のパイプライン次元を特定する。
• 最大パイプライン効率のための最適なトークン分割を計算する動的計画法ベースのアルゴリズムを開発する。

提案手法

• 単一入力シーケンス内のトークン次元にまたがってパイプライン処理を行うトークンレベルのパイプライン並列性を提案する。
• フォワード/バックワードの待機時間を、トークンスライスサイズとクラスターの特性の関数としてモデル化する。
• トークン次元上の最適なスライシング方式を見つけて訓練待機時間を最小化するための動的計画法アルゴリズムを開発する。
• 単純な性能モデルでフォワード伝搬時間を推定し、それをDP最適化の指針として用いる。
• 直交性を示す：TeraPipeは既存のデータ/モデル並列法（マイクロバッチによるパイプライン、演算分割、データ並列性）と組み合わせ可能である。

実験結果

リサーチクエスチョン

• RQ1 自己回帰Transformerにおいて、パイプライン並列性をレイヤー次元からトークン次元へ拡張するにはどうすればよいか？
• RQ2 特定のLMとクラスタに対して、トークン次元上のどのスライス方式が総訓練待機時間を最小化するか？
• RQ3 トークンレベルのパイプライン処理は、他のモデル並列技術やデータ並列性とどう相互作用するか？
• RQ4 トークンレベルのパイプライン並列性を用いて、大規模GPT-3クラスのモデルでどの程度の性能向上が達成できるか？
• RQ5 シーケンス長はトークンレベルのパイプライン並列性の有効性にどう影響するか？

主な発見

• TeraPipeは大規模LMに対して顕著な速度向上を生み出し、GPT-3-175Bの訓練を従来の同期的モデル並列法より最大5.0x速くする。48 AWS p3.16xlarge GPU上で。
• 動的計画法アプローチは最適なトークン分割スキームを効果的に決定し、パイプライン効率を最大化し、検討ケースで一様分割より約1.04x–1.12x上回る。
• メモリ制約によりバッチサイズが減少しパイプライン段数が増えるため、大きなモデルほど利得が大きく、トークンレベルのパイプライニングはより多くの飽和機会を提供する。
• 長い入力シーケンス長はトークンレベルパイプライン処理の潜在的利益を大幅に高め、観測される速度向上はシーケンス長の増加とともに増加する。
• TeraPipeは他の並列学習法と直交しており、マイクロバッチベースのパイプライン、演算分割、データ並列性など既存の並列訓練法と組み合わせ可能である。