サブロウ丸

Sabrou-mal サブロウ丸

主にプログラミングと数学

Megatron-LMのソースコードを読む

NVIDIAが提案するTransformerをベースとする言語処理モデルの並列化実装。サーベイ記事はこちら↓↓↓

Githubのレポジトリには

  • Data Preprocessing (データ前処理)
  • Pretraining(事前学習)
  • Evaluation and Tasks(評価)

のコードが含まれています。 事前学習については

  • BERT
  • GPT
  • T5

の3つのモデルについて、分散版とそうでないものがそれぞれ用意されていますね。 論文を読むと分かりやすいですが、Megatron-LMで提案しているのはモデルの垂直的な分割によるモデル並列とAttention部分の水平分割による並列処理の合わせ技です。

下記ではpretrain_bert.pyの実行を例にしてそれぞれを見ていきます。 ここではmodelの作成と訓練の大まかな流れが規定されています。 modelの作成はtraining::setup_model_and_optimizer > training::get_model で行われています。

モデルの並列化は torchDDP (torch.nn.parallel.distributed)とLocalDDP (megatron.model.DistributedDataParallel)が用意されています。前者はtorchがサポートしているデータ並列化のための機能。LocalDDPについては次項を参照。

並列実行コマンド

並列なし逐次実行と並列実行プログラムの差分。並列実行(+)の方は環境変数を指定しており、WORLD_SIZEがモデル並列の個数になります。

--- pretrain_bert.sh 2022-07-12 12:07:16.000000000 +0900
+++ pretrain_bert_distributed.sh  2022-06-06 16:29:40.000000000 +0900
@@ -1,32 +1,42 @@
 #!/bin/bash
 
-RANK=0
-WORLD_SIZE=1
+GPUS_PER_NODE=8
+# Change for multinode config
+MASTER_ADDR=localhost
+MASTER_PORT=6000
+NNODES=1
+NODE_RANK=0
+WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))
+
 DATA_PATH=<Specify path and file prefix>_text_sentence
 CHECKPOINT_PATH=<Specify path>
 
-python pretrain_bert.py \
+DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
+
+python -m torch.distributed.launch $DISTRIBUTED_ARGS \
+       pretrain_bert.py \
        --num-layers 24 \
        --hidden-size 1024 \
        --num-attention-heads 16 \
        --micro-batch-size 4 \
-       --global-batch-size 8 \
+       --global-batch-size 32 \
        --seq-length 512 \
        --max-position-embeddings 512 \
-       --train-iters 2000000 \
-       --lr-decay-iters 990000 \
+       --train-iters 1000000 \
        --save $CHECKPOINT_PATH \
        --load $CHECKPOINT_PATH \
        --data-path $DATA_PATH \
        --vocab-file bert-vocab.txt \
        --data-impl mmap \
        --split 949,50,1 \
+       --distributed-backend nccl \
        --lr 0.0001 \
-       --min-lr 0.00001 \
        --lr-decay-style linear \
-       --lr-warmup-fraction .01 \
+       --min-lr 1.0e-5 \
+       --lr-decay-iters 990000 \
        --weight-decay 1e-2 \
        --clip-grad 1.0 \
+       --lr-warmup-fraction .01 \
        --log-interval 100 \
        --save-interval 10000 \
        --eval-interval 1000 \

モデル分割

モデル並列を図にすると下記のようになっています。 実装においてはBERTではTransformerの構造を繰り返し積み上げるようなモデル設計がされています。その繰り返しのブロックについては入力と出力のサイズが同じであるため、特別な処理を行うことなくモデル構造の延長が可能になっています。実装を見るとtraining.py::get_modelでmodelのリストの中にBertModelを作ってappendしていますね。こんな感じでプログラム上では、あるモデルを分割して細分化するのではなく、初めから小さいモデルを複数用意してそれをつなげて一つのモデルを構成するようなボトルアップによる構成になっています。

パイプライン処理

モデル並列にはこのパイプライン処理はmegatron/schedules.pyで定義されている。たとえばforward_backward_pipelining_with_interleavingとか。training.py::train_stepの中ではこのように使われています。

具体的なコード

forward関数は何回もラップされているので見えずらくなっています。。 まず大きな分岐はmegatron/schedules.py::get_forward_backward_func()でパイプライン処理をするか否か、interleaveを有効にするかでforward, backward処理を分けます。 megatron/schedules.py::forward_backward_pipelining_with_interleaving()が複数デバイスによるモデル分割時に使用されるもので、この付近を見るように入力(input_tensors)を部分モデルに通して、出力をまた次の部分モデル用に更新します。input_tensors[model_chunk_id]がmodel_chunk_id番目の部分モデル用の入力。

LocalDDP

Megatron-LMではデータ並列モジュールとしてtorch.nn.parallel.distributedとLocalDDP(megatron自作)が用意されています。 LocalDDPはmegatron/model/distributed.pyでクラスとして実装されていれ、クラスメソッドとして下記を持つものです。

  • zero_grad_buffer
  • broadcast_params
  • allreduce_gradients
    • data-parallel中のgradientを集めている
    • どこで使っているか? > megatron/optimizer/optimizer.py::reduce_model_grads

all_reduce_gradientsが使用されているoptimizer.py::reduce_model_gradsは下記の内容(一部内容を削除)

def reduce_model_grads(self, args, timers):
    """All-reduce all grads, and all-reduce embeddings."""

    # All-reduce layer-norm grads (for sequence parallelism).
    self.allreduce_layernorm_grads(args)

    # All-reduce if needed.
    if args.DDP_impl == 'local':
        for model in self.models:
            model.allreduce_gradients()

    # All-reduce embedding grads.
    self.allreduce_embedding_grads(args)

pytorchのoptimizerはモデルパラメタの更新機能を行うものです。ここで

  1. allreduce_layernorm_grads
  2. allreduce_gradients
  3. allreduce_embedding_grads

の3つの勾配集約が行われていて、allreduce_layernorm_gradsとallreduce_embedding_gradsは共にoptimizer.pyの中で実装されている。

ParallelAttention

Attentinoは次のクラスで並列実装が行われています。

forwardではまずColumnParallelLinearでhidden_statesをもとに並列にQ, K, V行列を作成する。行列の形を成形したあとに、次にCoreAttentionでscaled-dot product attentionを行う。ただCoreAttentionでも入力する行列のshapeを変換して、隠れ双次元[h]を並列数に分割して次元を増やす[h//np, np]にしてtorchの行列演算(bmmなど)に通しているだけっぽい。(それで並列化できるからokってことか?)Attentionの部分はプロセス内並列のみを想定しているコードになっていますね。

ColumnParallelLinearの中のLinearWithGradAccumulationAndAsyncCommunicationでは行列積に関する同期処理が行われている。 docにはLinear layer execution with asynchronous communication and gradient accumulation fusion in backprop.とある。ちなみにforwardの引数にあるctxはcontextの略っぽい。