[論文レビュー] Sequence Parallelism: Making 4D Parallelism Possible

Shenggui Li, Fuzhao Xue+2 more|arXiv (Cornell University)|May 26, 2021

Advanced Neural Network Applications参考文献 23被引用数 12

ひとこと要約

この論文は、長時間のシーケンスを複数のGPUに分割することでメモリ効率の良い学習手法「シーケンス並列処理」を導入している。アテンション埋め込みをリング型アトムダウム（all-reduce）で通信することで、長時間のシーケンスモデリングが可能になる。64枚のP100 GPUで、トランスフォーマーのテンソル並列処理と比較して、バッチサイズが13.7倍、シーケンス長が3.0倍にまで拡大され、PyTorchにおけるデータ並列処理およびパイプライン並列処理とも完全に互換性がある。

ABSTRACT

Within Transformer, self-attention is the key module to learn powerful context-aware representations. However, self-attention suffers from quadratic memory requirements with respect to the sequence length, which limits us to process longer sequence on GPU. In this work, we propose sequence parallelism, a memory efficient parallelism method to help us break input sequence length limitation and train with longer sequence on GPUs. Compared with existing parallelism, our approach no longer requires a single device to hold the whole sequence. Specifically, we split the input sequence into multiple chunks and feed each chunk into its corresponding device (i.e. GPU). To compute the attention output, we communicate attention embeddings among GPUs. Inspired by ring all-reduce, we integrated ring-style communication with self-attention calculation and proposed Ring Self-Attention (RSA). Our implementation is fully based on PyTorch. Without extra compiler or library changes, our approach is compatible with data parallelism and pipeline parallelism. Experiments show that sequence parallelism performs well when scaling with batch size and sequence length. Compared with tensor parallelism, our approach achieved $13.7 imes$ and $3.0 imes$ maximum batch size and sequence length respectively when scaling up to 64 NVIDIA P100 GPUs. We plan to integrate our sequence parallelism with data, pipeline and tensor parallelism to further train large-scale models with 4D parallelism in our future work.

研究の動機と目的

• トランスフォーマーにおける自己アテンションの2次関数的メモリ増加を緩和し、GPU上のシーケンス長の制限を解消すること。
• アテンション計算中に、1つのデバイスが全入力シーケンスを保持する必要があるというボトルネックを克服すること。
• モデルやハードウェアの変更なしに、複数のGPUにシーケンスチャンクを分散させることで、長時間のシーケンス学習を効率的に行えるようにすること。
• 従来の並列処理戦略（データ並列処理やパイプライン並列処理など）と互換性を持つ、通信効率の良いアテンション機構を設計すること。
• 4D並列処理フレームワークにシーケンス並列処理を統合することで、バッチサイズとシーケンス長の両方をスケーラブルに拡大すること。

提案手法

• 入力シーケンスを複数のチャンクに分割し、それぞれを別々のGPUに割り当てることで、1つのデバイスが全シーケンスを保持する必要がなくなる。
• 自己アテンション計算中に、GPU間でアテンション埋め込みをリング型アトムダウムで通信することで、分散型アテンション計算を実現する。
• リング通信をアテンション計算パイプラインに直接統合した新規アテンション機構「リング自己アテンション（RSA）」を導入する。
• コンパイラやライブラリの変更なしに、PyTorchネイティブで実装することで、既存の学習フレームワークとの広範な互換性を確保する。
• データ並列処理およびパイプライン並列処理と組み合わせることで、大規模モデル学習に適した4D並列処理を実現する。
• 逆伝播処理中にリングアトムダウムを用いてデバイス間でアテンション出力を同期することで、勾配の一貫性とモデルの正確性を維持する。

実験結果

リサーチクエスチョン

• RQ1 複数のGPUにシーケンスチャンクを分散させることで、シーケンス並列処理はトランスフォーマーにおけるシーケンス長の制限を打ち破ることができるか？
• RQ2 64枚のP100 GPUでスケーリングした際、バッチサイズとシーケンス長のスケーリングにおいて、シーケンス並列処理はテンソル並列処理と比べてどの程度優れているか？
• RQ3 データ並列処理およびパイプライン並列処理と組み合わせた場合、シーケンス並列処理はバッチサイズとシーケンス長の両方をどの程度スケーラブルに拡大できるか？
• RQ4 アーキテクチャ的・コンパイラレベルの変更なしに、シーケンス並列処理をデータ並列処理やパイプライン並列処理といった既存の並列処理戦略と統合できるか？
• RQ5 標準的なアトムダウムやパラメータ共有と比較して、シーケンス並列処理におけるリング型通信は、アテンション計算においてどの程度効率的か？

主な発見

• 64枚のNVIDIA P100 GPUでスケーリングした際、シーケンス並列処理はテンソル並列処理と比較して、バッチサイズが13.7倍にまで拡大された。
• 同じハードウェア環境下で、シーケンス並列処理はテンソル並列処理と比較して、入力シーケンス長が3.0倍にまで拡張された。
• バッチサイズとシーケンス長の両方を増加させた場合でも、高いスケーラビリティを示し、良好な線形スケーリング効率を達成した。
• データ並列処理およびパイプライン並列処理と完全に互換性があり、既存の学習パイプラインへのシームレスな統合が可能である。
• リングアトムダウムを活用したリング自己アテンション（RSA）は、通信オーバーヘッドを効果的に低減し、学習中に高いスループットを維持した。
• PyTorchや外部ライブラリの変更が一切不要であり、既存のディープラーニングワークフローへの容易なデプロイが可能である。