Megatron-LM: モデル並列性を使用したマルチビリオンパラメータ言語モデルの学習

最近の言語モデルに関する研究では、大規模なトランスフォーマー・モデルの訓練が自然言語処理（Natural Language Processing, NLP）アプリケーションにおける最先端技術を進展させることを示しています。しかし、非常に大きなモデルはメモリ制約により訓練が困難であることがあります。本研究では、非常に大規模なトランスフォーマー・モデルの訓練手法を提示し、数十億のパラメータを持つトランスフォーマー・モデルの訓練を可能にする単層内並列化アプローチを実装します。当方針は新しいコンパイラやライブラリの変更を必要とせず、パイプライン並列化とは独立かつ補完的な関係にあり、ネイティブPyTorchにおいて数回の通信操作の挿入だけで完全に実装できます。我々はこのアプローチを512個のGPUを使用して最大83億パラメータを持つトランスフォーマー・ベースのモデルまで収束させる過程で説明します。強力な単一GPU基準（ピークFLOPsの30%に相当する39テラFLOPs）と比較して76%のスケーリング効率で15.1ペタFLOPsという全体的な性能を維持しました。大規模な言語モデルがさらなる最先端技術（State of the Art, SOTA）への進展をもたらすことを示すために、GPT-2に類似した83億パラメータのトランスフォーマー言語モデルとBERTに類似した39億パラメータのモデルを訓練しました。BERTのようなモデルにおいてレイヤー正規化（layer normalization）の配置に対する注意が、モデルサイズが増加するにつれて性能向上に不可欠であることを示しています。GPT-2モデルを使用してWikiText103データセット（SOTA困惑度15.8に対して10.8）とLAMBADAデータセット（SOTA精度63.2%に対して66.5%）で最先端結果を得ました。また、BERTモデルはRACEデータセット（SOTA精度89.4%に対して90.9%）で最先端結果を得ています。