[論文レビュー] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Mohammad Shoeybi, Mostofa Patwary+4 more|arXiv (Cornell University)|Sep 17, 2019

Topic Modeling参考文献 46被引用数 821

ひとこと要約

Megatron-LM は、PyTorch での層内モデル並列性を示し、数十億パラメータのトランスフォーマーモデルを訓練。512 GPUs で 8.3B パラメータを達成し、強いスケーリングと複数の NLP ベンチマークで最先端の結果を得ている。

ABSTRACT

Recent work in language modeling demonstrates that training large transformer models advances the state of the art in Natural Language Processing applications. However, very large models can be quite difficult to train due to memory constraints. In this work, we present our techniques for training very large transformer models and implement a simple, efficient intra-layer model parallel approach that enables training transformer models with billions of parameters. Our approach does not require a new compiler or library changes, is orthogonal and complimentary to pipeline model parallelism, and can be fully implemented with the insertion of a few communication operations in native PyTorch. We illustrate this approach by converging transformer based models up to 8.3 billion parameters using 512 GPUs. We sustain 15.1 PetaFLOPs across the entire application with 76% scaling efficiency when compared to a strong single GPU baseline that sustains 39 TeraFLOPs, which is 30% of peak FLOPs. To demonstrate that large language models can further advance the state of the art (SOTA), we train an 8.3 billion parameter transformer language model similar to GPT-2 and a 3.9 billion parameter model similar to BERT. We show that careful attention to the placement of layer normalization in BERT-like models is critical to achieving increased performance as the model size grows. Using the GPT-2 model we achieve SOTA results on the WikiText103 (10.8 compared to SOTA perplexity of 15.8) and LAMBADA (66.5% compared to SOTA accuracy of 63.2%) datasets. Our BERT model achieves SOTA results on the RACE dataset (90.9% compared to SOTA accuracy of 89.4%).

研究の動機と目的

• 単一 GPU のメモリ制限を超えて、数十億パラメータ級の言語モデルの訓練を動機づける。
• PyTorch に最小限の変更でフィットする、単純で効率的な層内モデル並列アプローチを開発する。
• GPT-2 および BERT様 architectures のスケーリング効率と性能を評価する。
• 言語モデリングと下流タスクで最先端の結果を示し、オープンソースコードを公開する。

提案手法

• 最小限の同期で MLP と自己注意の GEMM を GPU 間で分割し、層内モデル並列を実装する。
• 各トランスフォーマ層で前方伝播2回と後方伝播2回の全要約（all-reduce）通信を使用して、効率的なスケーリングを可能にする。
• 語彙次元に沿って入力埋め込み行列と出力埋め込み行列を並列化し、GPU間の通信を削減する。
• 追加の通信を避けるため、各 GPU で層正規化と残差計算を複製して維持する。
• 動的損失スケーリングと活性化チェックポイントを用いた混合精度 Transformer モデルをメモリ効率のために訓練する。

実験結果

リサーチクエスチョン

• RQ1 PyTorch における層内モデル並列はカスタムコンパイラなしでトランスフォーマーを数十億パラメータへスケールできるか。
• RQ2 モデルサイズが標準の NLP ベンチマークでの GPT-2 や BERT 風モデルの性能に、数十億パラメータへスケールした場合どう影響するか。
• RQ3 モデルサイズが大きくなるとき、性能を維持・向上させるために必要なアーキテクチャの調整（例: layer normalization の配置）は何か。
• RQ4 数百の GPU で数十億パラメータを訓練する際の実用的なスケーラビリティ限界（FLOPs、スループット、効率）はどこまでか。

主な発見

• 512 GPU で 8.3B パラメータを訓練し、8-way のモデル並列で最大 15.1 PetaFLOPs の連続実行と、強力な単一 GPU ベースライン（39 TeraFLOPs）に対して 76% のスケーリング効率を達成。
• モデル並列性はモデル単独およびモデル+データ構成の両方で強いウェークスケーリングをもたらし、最大の設定で 74%–77% のスケーリング効率を達成。
• BERT様モデルの層正規化の配置を慎重に行うことで、モデルサイズの増加に伴う性能の単調な向上を可能にする。
• GPT-2 スタイルのモデルは WikiText103 で最先端の perplexity（10.81）と強力な LAMBADA 精度（66.51%）を 8.3B パラメータで達成；RACE の精度もより大きな BERT/Megatron セットアップで 90.9% へ改善。
• BERT風モデルは最大 3.9B パラメータまで、提案アーキテクチャと訓練体制を用いればいくつかの GLUE風タスクと RACE で最先端の開発結果を達成。