NVIDIA が開発した大規模Transformerモデル向け分散学習フレームワーク。テンソル並列・パイプライン並列・シーケンス並列を統合的に提供し、数兆パラメータモデルの効率的な学習を実現する。GPT・BERT・T5等の主要アーキテクチャに対応する。
PC構成ビルダーで最適なパーツを選択
Megatron-LM は NVIDIA が2019年から開発する大規模Transformerモデル向けの分散学習フレームワークである。名称はTransformerの「変形」にちなみ、主にテンソル並列（Tensor Parallelism）とパイプライン並列（Pipeline Parallelism）を核とするモデル並列技術を提供する。DeepSpeedやFSDPがデータ並列の拡張を主軸とするのに対し、Megatron-LMはモデルの内部構造を分割する手法に特化している。
Megatron-LMの技術はNVIDIA NeMoフレームワークに統合されており、企業向けLLM学習プラットフォームとして広く利用されている。GPT-3（175B）、Bloom（176B）、Falcon（180B）など多くの大規模モデルの学習にMegatron-LMの技術が使われた。
テンソル並列はTransformerの個々のレイヤー内で重み行列をGPU間に分割する手法である。Megatron-LMはAttention機構とFeed-Forward Network（FFN）で異なる分割戦略を適用する。
AttentionのQKV射影行列を 列方向（Column Parallel） に分割する。N個のAttentionヘッドをN基のGPUに均等に割り当て、各GPUが独立にAttention計算を実行する。出力射影行列は 行方向（Row Parallel） に分割し、AllReduceで結果を統合する。
GPU 0: Q[:, :d/N], K[:, :d/N], V[:, :d/N] → Attention → O[:d/N, :] → AllReduce
GPU 1: Q[:, d/N:2d/N], K[:, d/N:2d/N], V[:, d/N:2d/N] → Attention → O[d/N:2d/N, :] → AllReduce
...
FFNの最初の線形層を列方向に分割し、GeLU活性化後の二番目の線形層を行方向に分割する。GeLU活性化が要素ごとの演算であるため、列分割後の各GPUで独立に実行でき、通信が不要となる。
| 演算 | 分割方向 | 通信 |
|---|---|---|
| QKV射影 | Column Parallel | なし |
| Attention計算 | Head並列 | なし |
| 出力射影 | Row Parallel | AllReduce（Forward） |
| FFN Layer 1 | Column Parallel | なし |
| GeLU活性化 | 各GPU独立 | なし |
| FFN Layer 2 | Row Parallel | AllReduce（Forward） |
1つのTransformerブロックあたりForwardで2回、Backwardで2回の合計4回のAllReduce通信が発生する。NVLink帯域（900GB/s）では数十マイクロ秒で完了するため、ノード内8GPUでは高い計算効率を維持できる。
パイプライン並列はモデルのレイヤー群を複数のステージに分割し、各ステージを異なるGPUに配置する。Megatron-LMは1F1B（One Forward One Backward）スケジューリングを採用し、バブル比率を最小化する。
ナイーブなパイプライン実行では、最初のステージがForwardを完了するまで後続ステージがアイドルとなり、Backward完了まで先頭ステージがアイドルとなる。P個のステージでマイクロバッチ数Mの場合、バブル比率は (P-1)/M で表される。
| スケジューリング | バブル比率 | メモリ使用量 | 特徴 |
|---|---|---|---|
| Naive（全Forward→全Backward） | (P-1)/M | M個分の活性値 | 単純だが非効率 |
| GPipe | (P-1)/M | M個分の活性値 | 全マイクロバッチForward後Backward |
| 1F1B | (P-1)/M | P個分の活性値 | ForwardとBackward交互実行 |
| Interleaved 1F1B | (P-1)/(M×V) | P×V個分の活性値 | 仮想ステージで細分化 |
1F1Bスケジューリングでは、ウォームアップフェーズで各ステージにマイクロバッチを流し、定常状態では1つのForwardと1つのBackwardを交互に実行する。活性値メモリがM個分からP個分に削減される利点がある。
Interleaved 1F1Bは各GPUに複数の仮想ステージ（V個）を割り当て、さらに細粒度のパイプラインを実現する。バブル比率が1/Vに改善されるが、通信回数が増加するトレードオフがある。
Megatron-LM v3で導入されたシーケンス並列は、LayerNorm・Dropout等のシーケンス次元方向の演算を分割する。テンソル並列と直交する並列化であり、テンソル並列ではカバーできない部分のメモリを削減する。
テンソル並列環境では、AllReduce後のLayerNormやDropout入力がGPU数分複製されている。シーケンス並列はこの複製をシーケンス次元で分割し、AllReduceをReduceScatter + AllGatherに置き換える。通信量は同等だが、活性値メモリがGPU数に比例して削減される。
Megatron-LM 0.7以降で導入されたコンテキスト並列は、長シーケンス学習向けにAttention計算のシーケンス次元をGPU間で分割する。Ring Attentionのバリアントを実装し、各GPUがKVキャッシュのチャンクを隣接GPUにリング状に送りながらAttention計算を進める。
| 並列化 | 分割次元 | 対象演算 | メモリ削減対象 |
|---|---|---|---|
| テンソル並列 | Head/Hidden次元 | QKV射影・FFN | 重みパラメータ |
| シーケンス並列 | シーケンス次元 | LayerNorm・Dropout | 活性値（非Attention部） |
| コンテキスト並列 | シーケンス次元 | Attention計算 | KVキャッシュ・Attention活性値 |
コンテキスト並列により、128Kトークン以上の超長シーケンスでもAttentionのメモリがGPU数に比例して削減される。
テンソル並列はノード内8GPU（NVLink接続）が前提のため、DGX/HGX構成のGPUクラスタで効果を発揮する。モデルサイズが7B以上でマルチノード学習する場合にMegatron-LMの採用を検討すべきである。7B未満ではFSDP/ZeROのみで十分なことが多い。
NeMoはMegatron-LMの技術をベースにNVIDIAが提供するエンタープライズ向けLLM開発プラットフォームである。Megatron-LMが分散学習エンジンを提供し、NeMoがデータ処理・モデル管理・評価・推論最適化を統合する。NeMo FrameworkからMegatron-LMの全機能にアクセスできる。
テンソル並列は各Transformerブロックで4回のAllReduce通信が発生し、帯域要求が非常に高い。NVLink（900GB/s）では効率が維持できるが、InfiniBand（50GB/s）では18倍遅くなりボトルネックとなる。物理的にNVLinkで接続されるのが通常8GPU（DGX H100構成）のため、TPは8GPUに限定される。16GPU NVSwitchシステムではTP=16も可能である。