> [!abstract] 概要 > 言語モデリングの最近の研究は、大規模な Transformer モデルの訓練が自然言語処理の性能を前進させることを示している。しかし非常に大きなモデルはメモリ制約により訓練が困難である。本論文では、非常に大規模な Transformer モデルを訓練する手法を提示し、数十億パラメータのモデル訓練を可能にする単純で効率的なレイヤー内モデル並列アプローチを実装する。我々のアプローチは新たなコンパイラやライブラリの変更を要せず、パイプラインモデル並列と直交・補完的であり、既存の PyTorch Transformer 実装へのわずかな通信演算の追加で完全に実装できる。このアプローチを 512 個の GPU を用いて最大 83 億パラメータの Transformer ベースモデルを収束させることで実証する。全体アプリケーションにわたり 15.1 PetaFLOPS を達成し、シングル GPU で 39 TeraFLOPS を維持する強力なベースライン比で 76% のスケーリング効率を達成する。大規模言語モデルが性能をさらに向上させることを実証するため、GPT-2 に類似した 83 億パラメータ Transformer 言語モデルと BERT に類似した 39 億パラメータモデルを訓練する。BERT 系モデルでは層正規化の配置に細心の注意を払うことがモデルサイズ増大に伴う精度向上に不可欠であることを示す。GPT-2 モデルで WikiText103(パープレキシティ 10.8、従来 SOTA の 15.8)と LAMBADA(正解率 66.5%、従来 SOTA の 63.2%)で SOTA を達成する。また BERT モデルで RACE データセットの SOTA(正解率 90.9%、従来 SOTA の 89.4%)を達成する。 ## 論文情報 | 項目 | 内容 | |---|---| | タイトル | Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism | | 著者 | Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro | | 所属 | NVIDIA | | 発表 | arXiv:1909.08053(2019 年 9 月投稿、2020 年 3 月 v4) | | URL | https://arxiv.org/abs/1909.08053 | | コード | https://github.com/NVIDIA/Megatron-LM | ## 概要 Megatron-LM は NVIDIA が 2019 年に発表した、Transformer モデル向けのレイヤー内テンソル並列(tensor parallelism)手法の基礎論文である。本論文は従来のパイプライン並列とは直交する「同一レイヤー内の行列演算を複数 GPU に分割する」設計を提案し、数十億パラメータの言語モデルを単純な PyTorch コードで分散訓練可能にした。Megatron-LM の名を冠する OSS フレームワークの起点であり、以後の 3D 並列化(tensor + pipeline + data)の礎となった基礎論文である。 ## 問題設定 GPT-2・BERT に代表される大規模言語モデルの訓練には、単一 GPU のメモリ容量を超えるパラメータとオプティマイザ状態の格納が必要となった。当時の先行手法には次の問題があった。 - **GPipe**: パイプライン並列を提供するが TensorFlow 専用コンパイラを要し、また「パイプラインバブル」と呼ばれる空き時間が発生する - **Mesh-TensorFlow**: 汎用的な分散テンソル計算を定義できるが、カスタムコンパイラに依存し既存実装の書き直しを要する 本論文の問いは「既存の PyTorch コードへの最小限の変更だけで、数十億パラメータの Transformer を単純かつ高効率に分散訓練できるか」である。 ## 提案手法 ### テンソル並列の基本原理 Megatron-LM のテンソル並列はレイヤー内分割(intra-layer model parallelism)と呼ばれ、1 つのレイヤーの行列演算を複数 GPU に分割する。Transformer の主要ブロック — MLP と自己注意 — それぞれに固有の分割法を適用する。 ### MLP ブロックの分割 MLP は 2 層の全結合(GEMM)から構成される。第一 GEMM の重み行列 $A$ を列方向に分割し $A = [A_1, A_2]$ とすることで、GeLU 非線形関数を各 GPU で独立に適用できる。 $[Y_1, Y_2] = [\text{GeLU}(XA_1), \text{GeLU}(XA_2)]$ これにより GeLU 前の同期点が不要になる。第二 GEMM は行方向に分割し、出力を GPU 間で削減(all-reduce)することで完結する。順伝播 1 回、逆伝播 1 回の計 2 回の all-reduce のみで、2 つの GEMM を GPU 群にまたがって実行できる。 実装は `f` と `g` という共役な演算子で表現される。`f` は順伝播で恒等写像、逆伝播で all-reduce を行い、`g` は順伝播で all-reduce、逆伝播で恒等写像を行う。 ### 自己注意ブロックの分割 多頭注意(multi-head attention)のアテンションヘッドは本質的に独立しているため、K/Q/V の GEMM をヘッド単位で列並列に分割できる。各 GPU は担当ヘッドの計算をローカルで完遂でき、ヘッド間の即時通信を要しない。後段の出力射影 GEMM は行並列で分割し、GPU 間で削減する。 この設計により Transformer の 1 レイヤーあたり順伝播 2 回・逆伝播 2 回、合計 4 回の all-reduce で全 GEMM を GPU 群に分散できる。 ### 語彙埋め込みの並列化 出力埋め込み行列 $E_{H \times v}$ は語彙次元に沿って列並列分割する。各 GPU が部分 GEMM $[XE_1, XE_2]$ を計算し、語彙サイズ起因の巨大な all-gather を避けるため、出力 GEMM とクロスエントロピー損失を融合して通信サイズを $b \times s \times v$ から $b \times s$ に削減する。 ### ハイブリッドモデル・データ並列化 テンソル並列はデータ並列と直交して組み合わせられる。同一ノード内の 2 〜 8 GPU がモデル並列グループを形成し、複数のモデル並列グループが並行してデータを処理するデータ並列グループを構成する。83 億パラメータモデルの場合、8-way モデル並列 × 64-way データ並列 = 512 GPU を使用する。 ### 層正規化の配置変更(BERT スケーリング) BERT 系モデルをスケールする際に元のアーキテクチャでは 336M パラメータを超えると精度が劣化することが確認された。解決策として層正規化と残差接続の順序を入れ替え、レイヤーへの入力前に層正規化を適用する(Pre-LayerNorm)方式を採用した。この変更により訓練の安定性が向上し、パラメータ数の増加に単調に精度が向上するようになった。 ### 実装の特徴 - 新たなコンパイラや C++ コードを要しない - PyTorch の `torch.autograd.Function` で数行のコードとして実装できる - DropOut の乱数生成をモデル並列領域の内外で異なるシードで管理する(内部はワーカーごとに独立なシード) ## 新規性 1. **カスタムコンパイラ不要の実装**: Mesh-TensorFlow が専用コンパイラを要するのに対し、既存 PyTorch への最小限の変更だけで実現 2. **MLP・注意の GEMM 融合で同期点削減**: MLP の列-行分割の組み合わせが GeLU 前の同期を除去し、1 レイヤー全体を 4 all-reduce で完結させる 3. **パイプライン並列との直交性**: 同一レイヤー内の分割のため、レイヤー間を分割するパイプライン並列と独立に組み合わせられる 4. **Pre-LayerNorm による大規模 BERT の安定化**: レイヤー正規化を入力前に置くことで 336M 超の BERT 系モデルの精度劣化を解消 ## 実験設定 - **ハードウェア**: 最大 32 台の NVIDIA DGX-2H サーバ(合計 512 個の Tesla V100 SXM3 32GB GPU)。ノード内帯域幅 300 GB/秒(NVSwitch)、ノード間帯域幅 100 GB/秒(8 × InfiniBand アダプタ) - **ベースライン**: 1.2B パラメータモデルをシングル GPU で実行、39 TeraFLOPS = 理論ピーク FLOPS の 30% - **スケーリング研究**: 1.2B / 2.5B / 4.2B / 8.3B の 4 構成を検証(隠れサイズ 1,536 〜 3,072、アテンションヘッドサイズを 96 で固定) - **GPT-2 評価**: 355M / 2.5B / 8.3B モデル、WikiText103 とLAMBADA でゼロショット評価 - **BERT 評価**: 336M / 1.3B / 3.9B モデル、MNLI / QQP / SQuAD 1.1 / SQuAD 2.0 / RACE でファインチューニング評価 - **訓練データ**: Wikipedia + CC-Stories + RealNews + OpenWebText から 174 GB の重複排除済みテキスト ## 実験結果 ### スケーリング効率 | パラメータ数 | GPU 数 | スケーリング効率 | |---|---|---| | 1.2B | 1 | 100%(ベースライン) | | 2.5B | 2 | 95% | | 4.2B | 4 | 82% | | 8.3B | 8 | 77% | モデル+データ並列の組み合わせでは最大 512 GPU で 74% のスケーリング効率を達成。全体の持続 FLOPs は 15.1 PetaFLOPS。 強スケーリングの測定(1.2B モデル固定、バッチサイズ固定):2 GPU で 1.64 倍、4 GPU で 2.34 倍、8 GPU で 2.98 倍の速度向上を確認。8 GPU を超えると per-GPU 計算量の減少と通信オーバーヘッドが支配的になる。 ### GPT-2 評価(ゼロショット) | モデル | WikiText103 パープレキシティ ↓ | LAMBADA 正解率 ↑ | |---|---|---| | 355M | 19.31 | 45.18% | | 2.5B | 12.76 | 61.73% | | 8.3B | 10.81 | 66.51% | | 従来 SOTA | 15.79 | 63.24% | 8.3B モデルが双方のベンチマークで SOTA を更新。モデルサイズ増大に伴い検証セットパープレキシティが単調に改善することも確認。 ### BERT 評価(ファインチューニング) 3.9B モデルが RACE テストセットで 90.9%(従来 SOTA の 89.4%)を達成。SQuAD 1.1 では F1 95.5/EM 90.0(アンサンブル)。336M モデルと 3.9B モデルを比較すると、すべての下流タスクでパラメータ増大による単調な精度向上を確認。 ## 考察 本論文は「コンパイラ不要」という実装の簡潔さを意図的に選択している点が重要である。Mesh-TensorFlow は汎用的に任意の並列化を記述できるが、専用コンパイラへの依存がエコシステムを制約する。Megatron-LM は Transformer のアーキテクチャ的な構造(多頭注意の独立性、MLP の 2 段 GEMM)を利用した「設計特化の簡単な実装」であり、汎用性を犠牲にして実用的な採用障壁を下げた。 テンソル並列の通信量は TP 並列度に正比例して増大するため、ノード内の高帯域インターコネクト(NVLink/NVSwitch)が前提となる。論文自体も「ノード内 8-way が実用上の限界」と示唆しており、ノード間の適用は通信律速になる。この制約は後続の 3D parallelism 研究が「TP をノード内、PP/DP をノード間」と構成する原則の根拠となった。 BERT の Pre-LayerNorm という副次的発見も重要で、大規模モデルの安定した訓練が単なるスケールアップではなく適切なアーキテクチャ設計と結合していることを示す。 ## 強み / 弱点・課題 **強み** - 実装の簡潔さ: PyTorch 数行の追加で完全な実装を実現 - パイプライン並列との直交性により 3D parallelism の礎を提供 - テンソル並列の通信最小化設計(1 レイヤー 4 all-reduce)が現代フレームワークでも標準的 - BERT Pre-LayerNorm が後続モデルアーキテクチャに広く採用 **弱点・課題** - テンソル並列はノード内(高帯域)が前提であり、ノード間適用は通信律速になる - 1-D 分割に限定しており、より効率的な 2-D / 2.5-D / 3-D テンソル並列(Optimus / Tesseract 等)は後続研究に委ねる - 16 GPU を超える大規模化にはパイプライン並列との組み合わせが必要で、本論文単独では扱っていない - 語彙埋め込みの並列化における GEMM サイズのパディング(語彙サイズを 1,024 の倍数に)は汎用ではない