# Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM > [!abstract] 概要 大規模言語モデルは多様なタスクで最高精度を達成しているが、効率的な訓練は困難である。理由は(a)GPU メモリ容量が限られており、大型モデルをマルチ GPU サーバ 1 台にも収められない場合があること、(b)計算量が膨大で非現実的に長い訓練時間を要することの 2 点である。その結果、テンソル並列やパイプライン並列といった新たなモデル並列手法が提案されてきた。しかし、これらの手法を素朴に使うと、数千 GPU 規模でのスケール問題が生じる。本論文では、テンソル並列、パイプライン並列、データ並列を組み合わせて数千 GPU にスケールする方法を示す。バブル比率を削減しつつ既存手法と同程度のメモリ消費量を保つ、新規のインターリーブドパイプラインスケジュールを提案する。本手法により、1 兆パラメータのモデルを 3072 GPU 上で 502 petaFLOP/s (GPU ごとの理論ピーク達成率 52%) で訓練イテレーションを実行できる。 ## 論文情報 - **著者**: Deepak Narayanan†★, Mohammad Shoeybi‡, Jared Casper‡, Patrick LeGresley‡, Mostofa Patwary‡, Vijay Korthikanti‡, Dmitri Vainbrand‡, Prethvi Kashinkunti‡, Julie Bernauer‡, Bryan Catanzaro‡, Amar Phanishayee∗, Matei Zaharia† - **所属**: ‡ NVIDIA、† Stanford University、∗ Microsoft Research - **会議**: SC '21 (International Conference for High Performance Computing, Networking, Storage, and Analysis), 2021年11月 - **DOI**: https://doi.org/10.1145/3458817.3476209 - **コード**: https://github.com/nvidia/megatron-lm ## 概要 本論文は **PTD-P**(Pipeline + Tensor + Data Parallelism)と呼ぶ 3D 並列化手法を提案し、以下の 3 つの技術革新を組み合わせる。 1. **インターリーブドパイプラインスケジュール** — バブル比率を (p-1)/m から (p-1)/(m·v) に削減。 2. **スキャッター・ギャザー通信最適化** — TP と PP の組み合わせを利用し、ノード間 InfiniBand の送受信量を 1/t に削減。 3. **融合カーネル** — バイアス+GeLU, バイアス+Dropout+Add のオペレータ融合と専用 Softmax カーネル。 3072 A100 GPU で 1 兆パラメータ GPT モデルを MFU 52%, 502 petaFLOP/s (集計) で訓練。この規模の訓練を実用的な時間(約 3 ヶ月)で行える最初の実証とされる。 ## 問題設定 LLM は指数的にパラメータが増加しており(GPT-3 175B 等)、2 つの根本的困難がある。 - **メモリ壁**: 80GB A100 GPUでもパラメータだけで 320GB 以上の大型モデルは収まらない。 - **計算時間**: GPT-3 の訓練を V100 1 枚で行うと約 288 年かかる試算。 従来の対処策: - **データ並列**: バッチをシャードするだけなのでモデルが 1 GPU に収まる必要がある。バッチサイズの下限と GPU 数の上限(=バッチサイズ)に制約。 - **テンソル並列(Megatron v1, 先行論文 [40])**: ノード内 NVLink で有効だが、ノード間(InfiniBand)に拡張すると通信コストが支配的になる。20B パラメータ以下では有効。 - **パイプライン並列(GPipe, PipeDream 等)**: ノード間に使えるが、パイプラインバブルがスループットを最大 50% 損なう。 **PTD-P** は 3 つを組み合わせてそれぞれの弱点を補う。 ## 提案手法 ### 2.1 データ並列 各ワーカーがモデル複製を持ち、勾配を定期的に集計して重みを同期。単体では大規模モデルに対応できないため、モデル並列の補完として使用。 ### 2.2 パイプライン並列 **デフォルトスケジュール (PipeDream-Flush / 1F1B)**: - ウォームアップフェーズでフォワードパスを p-1 個分実行し、その後フォワード→バックワードを交互に処理。 - バブル時間の割合: (p-1)/m ここで m = バッチ内のマイクロバッチ数、p = パイプラインステージ数。 - メモリは p 個以下のマイクロバッチのアクティベーションを保持するだけでよい(GPipe の m 個より大幅に少ない)。 **インターリーブドスケジュール(本論文の新提案)**: - 各デバイスが非連続な複数のレイヤー「チャンク」を担当する。例: 4 デバイスで 16 層なら、デバイス 1 が層 1-2, 9-10 を担当(v=2 のとき)。 - バブル比率: (p-1)/(m·v)。v 倍の削減。 - 代償: パイプライン通信量が v 倍になる(スキャッター・ギャザー最適化で補完)。 - 条件: m ≥ p (マイクロバッチ数がパイプラインステージ数の整数倍)。 ### 2.3 テンソル並列 Megatron-LM v1 の方式を踏襲し、トランスフォーマーの MLP と自己注意ブロックを GPU 間で分割: - **MLP ブロック**: A 行列を列方向に分割 → GeLU を各 GPU で独立適用 → B 行列を行方向に分割 → AllReduce で集約。フォワード/バックワードでそれぞれ 2 回の AllReduce。 - **自己注意ブロック**: Q/K/V 行列を列方向(ヘッド分割)で分割 → 出力射影を行方向で分割 → AllReduce 集約。 - f, g の共役演算子ペアで実装。f はフォワードで恒等写像、バックワードで AllReduce。g は逆。 ### 3. 性能解析: 並列化次元の選択原則 **基本方程式**: p × t × d = n (パイプライン × テンソル × データ × GPU 総数) **テイクアウェイ #1**: テンソル並列は g-GPU サーバ内(ノード内)に留める。それ以上はパイプライン並列を使う。 - 根拠: テンソル並列は AllReduce でノード間 InfiniBand を使うと通信コストが支配的になる。またテンソル次元を上げるほど行列計算が小さくなり GPU 利用率が低下する。 **テイクアウェイ #2**: モデル並列サイズ M = t × p でモデルがメモリに収まるように設定し、残りのリソースはデータ並列に割り当てる。 **テイクアウェイ #3**: ミクロバッチサイズ b はモデル・パイプライン深さ p・データ並列次数 d・バッチサイズ B に依存し、GPU 利用率とバブル削減のトレードオフを考慮して選ぶ。 ### 4. 実装 **通信最適化 (スキャッター・ギャザー)**: テンソル並列ランク間でアクティベーションが複製されるため、連続パイプラインステージへの送受信で同じテンソルが t 回送信される冗長性を解消。送信側でテンソルを t 個のチャンクに分割し各ランクが 1 チャンクを担当の InfiniBand カードで送信(スキャッター)、受信側で NVLink 上の AllGather で再構成(ギャザー)。通信量を 1/t に削減。 **計算最適化**: 1. データレイアウト変更: [b, s, a, h] → [s, b, a, h] に変更してメモリ集約的な転置を回避し、ストライドバッチ GEMM を使用。 2. 融合カーネル: PyTorch JIT でバイアス+GeLU と バイアス+Dropout+Add をオペレータ融合。 3. 専用 Softmax カーネル: scale + mask + softmax を融合したカーネル(GPT 用の因果マスク版と BERT 用の一般マスク版)。 ## 新規性 1. **3 種類の並列化の組み合わせの系統的分析**: TP/PP/DP の相互作用を解析的かつ実証的に初めて体系化した。 2. **インターリーブドパイプラインスケジュール**: バブル比率を v 倍削減しつつメモリ消費量をほぼ同等に保つ新スケジュール。先行する GPipe (全フォワード→全バックワード) や 1F1B スケジュールを改良。 3. **スキャッター・ギャザー通信最適化**: TP と PP を組み合わせた場合の冗長通信を自動的に除去する最適化。 4. **設計ガイドライン**: テンソル並列はノード内/パイプライン並列はノード間という最適配置原則を実証した。 ## 実験設定 - **ハードウェア**: NVIDIA Selene スーパーコンピュータ。各ノードに 80GB A100 GPU × 8、NVLink/NVSwitch 接続、200 Gbps HDR InfiniBand HCA × 8。3 階層ファットツリートポロジ(850 スイッチ)。 - **精度**: 混合精度(16 ビット)。A100 理論ピーク 312 teraFLOP/s。 - **モデル**: GPT アーキテクチャ。ボキャブラリサイズ 51,200、シーケンス長 2048。 - **規模**: 1.7B〜1.008T パラメータまでの 10 構成を評価。 - **比較対象**: ZeRO-3 (DeepSpeed)、非インターリーブドスケジュール。 ## 実験結果 ### 弱スケーリング (表1) | パラメータ数 | GPU 数 | MFU | 集計スループット | |---|---|---|---| | 1.7B | 32 | 44% | 4.4 PF/s | | 7.5B | 128 | 46% | 18.2 PF/s | | 145.6B | 1536 | 47% | 227.1 PF/s | | 310.1B | 1920 | 50% | 297.4 PF/s | | 529.6B | 2520 | 52% | 410.2 PF/s | | 1008.0B | 3072 | 52% | **502.0 PF/s** | 3072 GPU まで準線形スケーリングを達成(モデルが大きくなるほど行列サイズが大きくなり GPU 利用率が向上する超線形傾向)。 ### ZeRO-3 比較 (表2) - 175B GPT-3: PTD-P が ZeRO-3 より 6%(少 GPU 時)〜70%(GPU 数 2 倍時) 高スループット。 - 530B: PTD-P が ZeRO-3 より 24%(少 GPU 時)〜70%(GPU 数 2 倍時) 高スループット。 - スケール増加時の差が大きい理由: ZeRO-3 は GPU 数増加とともにノード間 AllReduce の通信量が増加するが、PTD-P はノード間通信を PP の安価な点対点通信に限定できる。 ### インターリーブドスケジュール (図12) - 175B GPT-3 (96 GPU) でインターリーブドスケジュールが非インターリーブドを最大 10% 上回るスループット。 - スキャッター・ギャザー最適化あり: 最大 11% のスループット向上 (530B で 133→148 teraFLOP/s per GPU)。 - スキャッター・ギャザー最適化なし: バッチサイズが大きいと非インターリーブドの方が優位になる。 ### 融合オペレータ (5.8節) - GPT-3 175B: 融合ありで 19% スループット向上 (113→135 teraFLOP/s per GPU)。 - 530B: 11% 向上 (133→148 teraFLOP/s per GPU)。 ### 通信帯域 (5.9節、1T パラメータモデル, 3072 GPU) - パイプライン並列通信 (点対点): 実効二分帯域幅 892 GB/s。 - データ並列通信 (AllReduce): 実効二分帯域幅 12.9 TB/s。 ### 訓練時間推定 式: 訓練時間 ≈ 8TP / (nX)。GPT-3 175B を 300B トークン訓練 → 1536 GPU で約 34 日。1T パラメータを 450B トークン訓練 → 3072 GPU で約 84 日。 ## 考察 PTD-P の成功は、3 並列化次元の通信コスト特性の差異を利用した配置にある。 - **TP のノード内閉込め**: AllReduce 通信を NVLink(ノード内高帯域)で完結させ、InfiniBand を不要にする。 - **PP のノード間使用**: 安価な点対点通信(per microbatch のアクティベーション転送 bsh のみ)でスケールアウト。 - **DP の外側**: バッチ単位の AllReduce で通信の頻度を最小化。 アクティベーション再計算は小バッチでは 33% のスループット低下をもたらすが、大バッチでは有効にする。パイプラインバブルが小さくなり全体スループットが最大 2 倍向上するケースも観測された。 ## 強み - **実証規模の大きさ**: 1T パラメータ, 3072 GPU という当時最大規模の実験。 - **系統的分析**: TP/PP/DP の相互作用を解析モデルと実験の両面で体系化した最初の論文。 - **オープンソース**: コードは github.com/NVIDIA/Megatron-LM で公開済み。 ## 弱点・課題 - **自動並列化配置の非対応**: FlexFlow, PipeDream, DAPPLE 等の自動探索と異なり、著者が提案するヒューリスティックに依存する。探索空間が複雑化すると最適解を見つけられない可能性がある。 - **GPT アーキテクチャに特化**: 評価はすべて GPT ベースのモデル。BERT など非均質なモデルアーキテクチャへの適用可能性は限定的(自己注意ブロックの分割方式が構造依存)。 - **同期パイプラインのみ**: PipeDream-2BW, PipeMare 等の非同期/有界非同期方式は除外。これらはスループットで優位な場合があるが収束精度のトレードオフがある。 - **モデルサイズ拡大への依存**: MFU が 1.7B(44%)から 1T(52%)へ改善する背景には、モデルが大きくなると行列乗算サイズが増え GPU 利用率が上がるという性質がある。小モデルでは同等の効率を出しにくい。