Merge branch 'pipeline_parallel_main' into 'main'

[Megatron](https://arxiv.org/pdf/1909.08053.pdf) is a large, powerful transformer developed by the Applied Deep Learning Research team at NVIDIA. This repository is for ongoing research on training large transformer language models at scale. We developed efficient, model-parallel, and multinode training of [GPT-2](https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) and [BERT](https://arxiv.org/pdf/1810.04805.pdf) using mixed precision.

[Megatron](https://arxiv.org/pdf/1909.08053.pdf) is a large, powerful transformer developed by the Applied Deep Learning Research team at NVIDIA. This repository is for ongoing research on training large transformer language models at scale. We developed efficient, model-parallel (tensor and pipeline), and multinode training of [GPT-2](https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) and [BERT](https://arxiv.org/pdf/1810.04805.pdf) using mixed precision.

Using our GPT-2 model we achieve a perplexity of 10.8 on the WikiText-103 dataset (improving SOTA from 15.8) and an accuracy of 66.5% on the LAMBADA datasets. For BERT training, we swapped the position of the layer normalization and the residual connection in the model architecture (similar to GPT-2 architucture), which allowed the models to continue to improve as they were scaled up. Our BERT model with 3.9 billion parameters reaches a loss of 1.16, SQuAD 2.0 F1-score of 91.7, and RACE accuracy of 90.9%.

The two tiers of parallelism are data and model parallelism. First, we facilitate two distributed data parallel implementations: a simple one of our own that performs gradient all-reduce at the end of back propagation step, and Torch's distributed data parallel wrapper that overlaps gradient reduction with back propagation computation. To switch between these two options use `--DDP-impl local` or `--DDP-impl torch`, respectively. As expected, Torch distributed data parallelism is more efficient at larger model parallel sizes. For example, for the 8.3 billion parameters model running on 512 GPUs, the scaling increases from 60% to 76% when Torch's distributed data parallel is used. However, the overlapping method requires more memory and for some configurations (e.g., 2.5 billion parameters using 2-way model parallel and 1.2 billion parameters with no model parallel) can make the overall training slower as a result. We empirically found that using a smaller model in those cases improves the training time.

Second, we developed a simple and efficient intra-layer model parallel approach. To use model parallelism, add the `--model-parallel-size` flag to specify the number of GPUs among which to split the model, along with the arguments passed to the distributed launcher as mentioned above. With `WORLD_SIZE` GPUs and `MP_SIZE` model parallel size, `WORLD_SIZE`/`MP_SIZE` GPUs will be used for data parallelism. The default value for `--model-parallel-size` is 1, which will not implement model parallelism.

Second, we developed a simple and efficient two-dimensional model-parallel approach. To use tensor model parallelism (splitting execution of a single transformer module over multiple GPUs), add the `--tensor-model-parallel-size` flag to specify the number of GPUs among which to split the model, along with the arguments passed to the distributed launcher as mentioned above. To use pipeline model parallelism (sharding the transformer modules into stages with an equal number of transformer modules on each stage, and then pipelining execution by breaking the batch into smaller microbatches), use the `--pipeline-model-parallel-size` flag to specify the number of stages to split the model into (e.g., splitting a model with 24 transformer layers across 4 stages would mean each stage gets 6 transformer layers each). The number of microbatches in a per-pipeline minibatch is controlled by the `--num-microbatches-in-minibatch` argument. With `WORLD_SIZE` GPUs, `TENSOR_MP_SIZE` tensor-model-parallel size, `PIPELINE_MP_SIZE` pipeline-model-parallel-size, `WORLD_SIZE`/(`TENSOR_MP_SIZE` * `PIPELINE_MP_SIZE`) GPUs will be used for data parallelism. The default values for `--tensor-model-parallel-size` and `--pipeline-model-parallel-size` is 1, which will not implement either form of model parallelism.

We have examples of how to use these two different forms of model parallelism in these scripts:

`bash examples/pretrain_bert_distributed_with_mp.sh`

`bash examples/pretrain_gpt2_distributed_with_mp.sh`

Other than these minor changes, the distributed training is identical to the training on a single GPU.

                --save $CHECKPOINT_PATH \

                --load $CHECKPOINT_PATH \

                --data-path $DATA_PATH \

                --model-parallel-size $MP_SIZE \

                --tensor-model-parallel-size $MP_SIZE \

                --DDP-impl torch

</pre>

                --save $CHECKPOINT_PATH \

                --load $CHECKPOINT_PATH \

                --data-path $DATA_PATH \

                --model-parallel-size $MP_SIZE \

                --tensor-model-parallel-size $MP_SIZE \

                --DDP-impl torch

</pre>

Because evaluation requires substantially less memory than training, it may be advantageous to merge a model trained in parallel for use on a single GPU in downstream tasks. The following script accomplishes this.

<pre>

MODEL_PARALLEL_SIZE=2

TENSOR_MODEL_PARALLEL_SIZE=2

VOCAB_FILE=bert-vocab.txt

CHECKPOINT_PATH=checkpoints/bert_345m

WORLD_SIZE=$MODEL_PARALLEL_SIZE python tools/merge_mp_partitions.py \

WORLD_SIZE=$TENSOR_MODEL_PARALLEL_SIZE python tools/merge_mp_partitions.py \

        --model-type BERT \

        --model-parallel-size $MODEL_PARALLEL_SIZE \

        --tensor-model-parallel-size $TENSOR_MODEL_PARALLEL_SIZE \

        --tokenizer-type BertWordPieceLowerCase \

        --vocab-file $VOCAB_FILE \

        --num-layers 24 \

## BERT Task Evaluation

<a id="race-evaluation"></a>

### RACE Evaluation

The following script finetunes the BERT model for evaluation on the [RACE dataset](http://www.cs.cmu.edu/~glai1/data/race/). The `TRAIN_DATA` and `VALID_DATA` directory contain the RACE dataset as separate `.txt` files.

The following script finetunes the BERT model for evaluation on the [RACE dataset](http://www.cs.cmu.edu/~glai1/data/race/). The `TRAIN_DATA` and `VALID_DATA` directory contain the RACE dataset as separate `.txt` files. Note that for RACE, the batch size is the number of RACE query's to evaluate. Since each RACE query has four samples, the effective batch size passed through the model will be four times the batch size specified on the command line.

<pre>

TRAIN_DATA="data/RACE/train/middle"

               --vocab-file $VOCAB_FILE \

               --merge-file $MERGE_FILE \

               --load $CHECKPOINT \

               --model-parallel-size 1 \

               --tensor-model-parallel-size 1 \

               --num-layers 24 \

               --hidden-size 1024 \

               --num-attention-heads 16 \

               --vocab-file $VOCAB_FILE \

               --epochs 5 \

               --pretrained-checkpoint $PRETRAINED_CHECKPOINT \

               --model-parallel-size 1 \

               --tensor-model-parallel-size 1 \

               --num-layers 24 \

               --hidden-size 1024 \

               --num-attention-heads 16 \

               --vocab-file $VOCAB_FILE \

               --epochs 3 \

               --pretrained-checkpoint $PRETRAINED_CHECKPOINT \

               --model-parallel-size 1 \

               --tensor-model-parallel-size 1 \

               --num-layers 24 \

               --hidden-size 1024 \

               --num-attention-heads 16 \

MERGE_FILE=gpt2-merges.txt

python tools/generate_samples_gpt2.py \

       --model-parallel-size 1 \

       --tensor-model-parallel-size 1 \

       --num-layers 24 \

       --hidden-size 1024 \

       --load $CHECKPOINT_PATH \