Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

arxiv.org

最大1 trillionパラメータのtransformerを3000GPUs規模で分散学習する方法について模索するNVIDIAの論文

方法

• data parallelism……各ワーカーが同じモデルを持ち違うデータに対してfwd/bwdをして勾配を計算し、勾配をallreduceする
• model parallelism……各ワーカーがモデルの別の部分を担当する
  • tensor parallelism……各ステップの計算を分担する。例えばmulti-head attentionであれば単に各headを違うワーカーで計算する。
  • pipeline parallelism……モデルの一連の計算を部分部分に分割し担当する。mini-batchをmicro-batchに分割し、異なるmicro-batchの計算をoverlapさせる。GoogleのGPipeが有名。

その他、activation recomputation（gradient checkpointing）なども使う。

実装

github.com