[論文レビュー] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

Deepak Narayanan, Mohammad Shoeybi+10 more|arXiv (Cornell University)|Apr 9, 2021

Topic Modeling被引用数 7

ひとこと要約

本稿では、3072個のGPU上でトランジリオンパラメータ言語モデルの効率的学習を可能にする、メガトロン-LMにおけるパイプライン、テンソル、データ並列性の新規組み合わせであるPTD-Pを提案する。インタリーブドパイプラインスケジューリングを導入することで、502ペタFLOP/sのスループットを達成し、理論ピークの52％のGPU利用率を実現した。これは、メモリ効率を維持したままスループットを顕著に向上させた。

ABSTRACT

Large language models have led to state-of-the-art accuracies across a range of tasks. However, training these models efficiently is challenging for two reasons: a) GPU memory capacity is limited, making it impossible to fit large models on even a multi-GPU server, and b) the number of compute operations required to train these models can result in unrealistically long training times. Consequently, new methods of model parallelism such as tensor and pipeline parallelism have been proposed. Unfortunately, naive usage of these methods leads to fundamental scaling issues at thousands of GPUs, e.g., due to expensive cross-node communication or devices spending significant time waiting on other devices to make progress. In this paper, we show how different types of parallelism methods (tensor, pipeline, and data parallelism) can be composed to scale to thousands of GPUs and models with trillions of parameters. We survey techniques for pipeline parallelism and propose a novel interleaved pipeline parallelism schedule that can improve throughput by 10+% with memory footprint comparable to existing approaches. We quantitatively study the trade-offs between tensor, pipeline, and data parallelism, and provide intuition as to how to configure distributed training of a large model. Our approach allows us to perform training iterations on a model with 1 trillion parameters at 502 petaFLOP/s on 3072 GPUs with achieved per-GPU throughput of 52% of theoretical peak. Our code is open sourced at https://github.com/nvidia/megatron-lm.

研究の動機と目的

• GPUメモリ容量を超える大規模言語モデルの学習を、実行に耐えられない長時間に及ぶ課題に対処するため。
• 数千のGPUを横断して、1兆パラメータを超えるモデルのスケーラブルな学習を可能にするため。
• 最適化子の厳密な意味論を保ちながら、パイプライン、テンソル、データ並列性の組み合わせを最適化するため。
• 新しいスケジューリング戦略により、パイプラインフラッシュのオーバーヘッドを低減し、小バッチサイズでのスループットを向上させるため。
• 大規模GPUクラスタ上で高い計算効率（理論ピークの52％）を達成するため。

提案手法

• ノード間のパイプライン並列性、ノード内でのテンソル並列性、およびデータ並列性（PTD-P）を組み合わせ、数千のGPUにわたる学習スケーリングを実現する。
• パイプラインフラッシュ時間を短縮し、ベースライン手法と比較して10％以上のスループット向上を実現する、インタリーブドパイプラインスケジューリングを導入する。
• アクティベーション再計算を用いてメモリ使用量を削減し、GPUメモリ制限内でより大きなモデルを収容可能にする。
• オペレータの統合と最適化されたデータレイアウトを適用して、メモリバウンドカーネルと通信オーバーヘッドを最小限に抑える。
• モデルのデバイス間分割をガイドするコストモデルを活用し、通信量を最小限に抑えつつデバイスの利用率を最大化する。
• サーバ内での高帯域幅インタコネクト（例：NVLink）と最適化されたノード間通信を活用して、マルチノード学習における遅延を低減する。

実験結果

リサーチクエスチョン

• RQ1 パイプライン、テンソル、データ並列性を効果的に組み合わせることで、数千のGPUにわたる学習をスケーリングする方法は何か？
• RQ2 パイプラインフラッシュのオーバーヘッドを最小限に抑え、特に小バッチサイズでスループットを最大化するスケジューリング戦略は何か？
• RQ3 GPUクラスタ上でトランジリオンパラメータモデルを学習する際のメモリ制約をどのように克服できるか？
• RQ4 アクティベーション再計算と通信最適化が学習効率に与える影響は何か？
• RQ5 提案手法が、大規模GPUクラスタ上で高い利用率（例：ピークの52％）をどの程度達成できるか？

主な発見

• 提案されたインタリーブドパイプラインスケジューリングにより、標準的なパイプライン手法と比較して、学習スループットが10％以上向上した。
• 3072個のGPU上で、ピークの502ペタFLOP/sの学習スループットを達成し、理論GPUピーク性能の52％に到達した。
• 提案されたPTD-Pフレームワークを用いることで、1兆パラメータのモデルの学習が約3か月で可能になった。
• テンソル、パイプライン、データ並列性の組み合わせにより、単一またはマルチGPUサーバーのメモリおよび計算限界を超えて、スケーラブルな学習が可能になった。
• アクティベーション再計算により、FLOP数の著しい増加を伴わずにメモリ使用量が削減され、より大きなモデルの学習が可能になった。
• パイプライン同期に起因する無駄な待機時間を最小限に抑えつつ、最適化子の厳密な意味論を維持した。