- Transformer層内の行列を複数GPUへ分割し、大きいモデルを学習した。
- 83億パラメータの構成を512 GPUで訓練し、アプリケーション全体で最大15.1 PFLOPSを報告した。
- 76%の拡張効率は特定の比較構成・設備に対する値だ。
何を確かめた研究か
モデルのパラメータや中間計算が単一GPUに収まらない時、通常のデータ並列だけでは解決できない。Megatron-LMはTransformerの注意計算と多層パーセプトロンを層内で分割し、必要な同期を限定して複数GPUへ載せる方法を調べた。目的は単なる最大パラメータ数ではなく、訓練を実際に進められる処理効率だ。
どう検証したか
行列の列方向・行方向の分割を組み合わせ、各GPUでできる計算を進めてから必要な結果を集約する。論文はモデル並列を8 GPU、データ並列を64群まで組み合わせ、合計512 GPUの構成を使った。12億から83億パラメータまでのモデルで、計算量、通信、訓練の収束、言語モデルの予測やLAMBADAを評価した。基準の単一GPU構成も明示する。
何が分かったか
要旨と§5では、83億パラメータのモデルを512 GPUで動かし、全アプリケーションで最大15.1ペタFLOPS、単一GPU基準に対する拡張効率76%を報告する。表1の構成では8 GPUのモデル並列と64群のデータ並列である。15.1という数値はサービスの推論速度ではなく、訓練計算の処理量だ。
複数GPUに分けると通信が必要になり、GPU数を増やすだけで線形に速くなるわけではない。論文は同一サイズのモデルをより多いGPUへ分割する強いスケーリングも測り、モデルを大きくする実験と区別している。モデルの性能評価は並列手法の成立を支えるが、モデル規模だけが品質向上の原因と証明する設計ではない。
限界は何か
設備は当時のDGX-2H、Tesla V100、GPU間の高速通信に依存する。異なるクラスタや通信速度では効率が変わる。パラメータ数と訓練FLOPSだけでは、費用や電力、実際の利用者への価値は分からない。論文は言語モデル訓練を主に扱い、推論の低遅延を直接示していない。
仕事との接点
大きいモデルの運用を検討する場合は、訓練時と推論時の負荷を分けて予算化する。GPUの枚数だけでなく、通信待ち、実行時間、対象タスクの改善を測り、並列化の費用に見合うか確認する。