## Memo ## Memo with LLM ### 論文情報 - **論文のタイトル**: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM - **著者と所属**: - Deepak Narayanan (Stanford University, NVIDIA) - Mohammad Shoeybi (NVIDIA) - Jared Casper (NVIDIA) - Patrick LeGresley (NVIDIA) - Mostofa Patwary (NVIDIA) - Vijay Korthikanti (NVIDIA) - Dmitri Vainbrand (NVIDIA) - Prethvi Kashinkunti (NVIDIA) - Julie Bernauer (NVIDIA) - Bryan Catanzaro (NVIDIA) - Amar Phanishayee (Microsoft Research) - Matei Zaharia (Stanford University) - **カンファレンス/ジャーナル名**: SC '21: International Conference for High Performance Computing, Networking, Storage and Analysis - **発表年**: 2021 ### 論文概要 この論文は、GPU クラスターでの大規模言語モデル訓練を効率的に行うためのパラレル化戦略を提案している。テンソルパラレリズム、パイプラインパラレリズム、データパラレリズムを組み合わせた PTD-P という手法により、3072個のA100 GPUで1兆パラメータのモデルを502 petaFLOP/sの処理能力で訓練することを実現した。 ### 詳細解説 #### 問題設定 **入力**: 数兆パラメータを持つ大規模な[[Transformer]]ベース言語モデル（GPTなど）の訓練 **出力**: 効率的な分散訓練により訓練されたモデル **必要なデータ**: - テキストデータセット（語彙サイズV=51,200、シーケンス長s=2048） - 大規模GPUクラスター（最大3072個のNVIDIA A100 GPU） - 高速ネットワーク（InfiniBandによる相互接続） 従来の課題として、(a)単一GPUのメモリ容量では大規模モデルが収まらない、(b)計算量が膨大で現実的でない訓練時間が必要、という2つの問題があった。 #### 提案手法 **PTD-P（Pipeline-Tensor-Data Parallelism）**という3つのパラレリズムを組み合わせた手法を提案： 1. **テンソルパラレリズム（Tensor Parallelism）**: - Transformerの各層内でGEMMを複数GPUに分割 - MLPブロック：重み行列Aを列方向に分割 A = [A₁, A₂] - [[Self Attention]]：Key、Query、Value行列を列方向に分割 - フォワード・バックワードパスでそれぞれ2回のall-reduce通信が必要 2. **パイプラインパラレリズム（Pipeline Parallelism）**: - モデル層を複数デバイスに分散配置 - バッチを複数のマイクロバッチに分割してパイプライン実行 - 新規スケジュール「Interleaved 1F1B」を提案：各デバイスが複数のモデルチャンクを担当することでパイプラインバブルを1/v倍に削減 3. **データパラレリズム（Data Parallelism）**: - 入力データセットを分割し、複数ワーカーで勾配を集約 **バブル時間削減の数式**: - 従来スケジュール: (p-1)/m - インターリーブスケジュール: (1/v)·(p-1)/m （p: パイプライン段数、m: マイクロバッチ数、v: チャンク数） **最適化技術**: - Scatter/Gather通信最適化：テンソル並列ランク間で同じデータを8回送信する冗長性を排除 - 計算最適化：データレイアウト変更、カーネル融合（bias+GeLU、scale+mask+softmax） #### 新規性 1. **3種のパラレリズムの初の体系的組み合わせ**: 既存手法は単一または2種の組み合わせのみ 2. **新規パイプラインスケジュール**: Interleaved 1F1Bスケジュールによりバブル時間を大幅削減 3. **通信最適化**: Scatter/Gather手法でノード間通信を1/t倍に削減 4. **パフォーマンス解析**: 異なるパラレリズム次元の相互作用を定量的に分析 先行研究との比較： - ZeRO-3: モデル並列なしのデータ並列手法と比較して最大70%の性能向上 - PipeDream, GPipe: パイプライン並列のみ - Megatron (原版): テンソル並列のみで20B パラメータが限界 #### 実験設定 **使用データセット**: GPTモデル（1B〜1T パラメータ） **評価指標**: - GPU当たりのteraFLOP/s - 理論ピーク性能に対する割合 - 総処理能力（petaFLOP/s） **実験環境**: - NVIDIA Selene スーパーコンピュータ - 各ノード8個のNVIDIA A100 80GB GPU - NVLink/NVSwitch内部接続 - Mellanox 200Gbps HDR InfiniBand **モデル設定例（1T パラメータ）**: - アテンションヘッド数: 160 - 隠れ次元: 25,600 - 層数: 128 - テンソル並列度: 8 - パイプライン並列度: 64 #### 実験結果 **主要な性能結果**: - **1Tパラメータモデル**: 3072 GPU で 163 teraFLOP/s/GPU（理論ピークの52%）、総計502 petaFLOP/s - **175B パラメータモデル（GPT-3相当）**: 1024 GPU で推定訓練時間34日 - **1T パラメータモデル**: 3072 GPU で推定訓練時間84日 **スケーリング結果**: - 32 GPU（1.7B パラメータ）から3072 GPU（1T パラメータ）まで超線形スケーリングを達成 - GPU利用効率が44%から52%に向上（大きな行列乗算でGPU利用率向上） **比較結果**: - ZeRO-3との比較: 175Bおよび530Bパラメータモデルで最大70%の性能向上 - ノード間通信量削減: scatter/gather最適化により最大11%の性能向上 **各最適化の効果**: - カーネル融合: GPT-3で19%、530Bモデルで11%の性能向上 - インターリーブスケジュール: 最大10%の性能向上 - 活性化再計算: メモリ効率的だが最大33%の性能低下、ただし大バッチサイズでは2倍の性能向上 ## Abstract 大規模言語モデルは、様々なタスクにわたって最先端の精度を実現している。しかし、これらのモデルを効率的に訓練することは、次の2つの理由から困難である：a) GPUメモリ容量が限られており、大規模モデルをマルチGPUサーバーでも収容することが不可能である、b) これらのモデルの訓練に必要な計算操作数が、非現実的に長い訓練時間をもたらす可能性がある。その結果、テンソル並列やパイプライン並列などの新しいモデル並列手法が提案されている。しかし残念ながら、これらの手法の単純な使用は、数千のGPUにおいて根本的なスケーリングの問題を引き起こす。例えば、高価なノード間通信や、他のデバイスの進行を待つことに多くの時間を費やすデバイスなどが原因である。本論文では、テンソル、パイプライン、データ並列の異なる種類の並列手法をどのように組み合わせて数千のGPUや兆のパラメータを持つモデルにスケールできるかを示す。我々は、既存のアプローチと同等のメモリフットプリントで10%以上のスループット向上を実現できる新しいインターリーブパイプライン並列スケジュールを提案する。我々のアプローチにより、3072個のGPUで1兆パラメータのモデルに対して502ペタFLOP/sで訓練イテレーションを実行することが可能となった（GPUあたりのスループットは理論ピークの52%）。