Megatron是由 NVIDIA 应用深度学习研究团队开发的大型、强大的变压器。该存储库用于正在进行的大规模培训大型转换器语言模型的研究。我们使用混合精度开发了GPT和BERT的高效、模型并行（张量和管道）和多节点预训练。

以下是我们直接使用威震天的一些项目：

• 使用 Megatron 进行 BERT 和 GPT 研究
• BioMegatron：更大的生物医学领域语言模型
• 用于开放域问答的神经检索器的端到端训练
• 大规模多角色生成对话建模
• 本地知识驱动的会话代理
• MEGATRON-CNTRL：使用大规模语言模型通过外部知识生成可控的故事
• RACE 阅读理解数据集排行榜
• 从合成数据训练问答模型

我们的代码库能够有效地训练具有模型和数据并行性的非常大（数千亿个参数）的语言模型。为了演示代码如何随着多个 GPU 和模型大小进行扩展，我们考虑了从 10 亿到 1 万亿参数的 GPT 模型。所有模型都使用 51,200 的词汇量和 2048 的序列长度。我们改变隐藏大小、注意力头数和层数以达到特定的模型大小。随着模型大小的增加，我们也会适度增加批量大小。我们利用NVIDIA 的 Selene 超级计算机来执行扩展研究并使用多达 3072 个A100用于最大型号的 GPU。下表显示了模型配置以及每秒实现的 FLOP（每个 GPU 和所有 GPU 的聚合）。请注意，FLOP 是针对端到端训练测量的，即包括所有操作，包括数据加载、优化甚至日志记录。

下图显示了达到的理论峰值 FLOP 百分比和每秒达到的聚合 petaFLOP 与 GPU 数量的关系。从 10 亿到 1 万亿的所有案例都实现了 41% 以上的半精度利用率，这对于端到端的应用来说是很高的。我们观察到，最初随着模型并行大小的增加，利用率略有下降；随着较大模型的隐藏大小增加，利用率开始增加，最大模型的利用率达到 49%。我们还注意到，在所有 GPU 上实现的每秒聚合 petaFLOPs 几乎随 GPU 数量线性增加，这表明良好的弱扩展性。

内容

• 内容
• 设置
  • 下载检查点
• 用法
• 训练
  • 数据预处理
  • BERT 预训练
  • GPT 预训练
  • 分布式 BERT 或 GPT 预训练
  • GPT-3 示例
• 评估和任务
  • GPT 文本生成
  • GPT 评估
    • WikiText 困惑度评估
    • LAMBADA 完形填空精度
  • BERT 任务评估
    • 种族评估
    • MNLI 评估
• 数据集
  • 收集维基百科训练数据
  • 收集 GPT Web 文本数据

设置

我们使用NGC 的 PyTorch 容器版本 20.12 测试了威震天，该容器使用 python 3.8、pytorch 1.8、cuda 11.1 和 nccl 2.8.3。

要使用此存储库，请安装支持 GPU 的最新支持版本的 PyTorch（python 3.8、pytorch 1.8、cuda 11.1 和 nccl 2.8.3 及更高版本）和 NVIDIA APEX。我们强烈建议使用NGC 最近的 PyTorch 容器之一（发布时最新的兼容版本可以使用 提取docker pull nvcr.io/nvidia/pytorch:20.12-py3）。数据预处理需要NLTK，但训练、评估或下游任务不需要。

下载检查点

我们提供了预训练的 BERT- 345M和GPT-345M检查点，用于评估或微调下游任务。要访问这些检查点，请先注册并设置NVIDIA GPU Cloud (NGC) Registry CLI。下载模型的更多文档可以在NGC 文档中找到。

或者，您可以使用以下命令直接下载检查点：

BERT-345M-uncased: wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.1_cased/zip -O megatron_bert_345m_v0.1_uncased.zip
BERT-345M-cased: wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.1_cased/zip -O megatron_bert_345m_v0.1_cased.zip
GPT-345M: wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/megatron_lm_345m/versions/v0.0/zip -O megatron_lm_345m_v0.0.zip

模型需要词汇文件才能运行。BERT WordPiece vocab 文件可以从 Google 的预训练 BERT 模型中提取：uncased、cased。GPT词汇文件和合并表可以直接下载。

用法

安装后，有几种可能的工作流程。最全面的是：

1. 数据预处理
2. 预训练
3. 微调（零样本任务可选）
4. 下游任务评估或文本生成

但是，步骤 1 和 2 可以通过使用上述预训练模型之一来替换。

我们提供了几个用于在examples目录中预训练 BERT 和 GPT 的脚本，以及用于零样本和微调下游任务的脚本，包括 MNLI、RACE、WikiText103 和 LAMBADA 评估。还有一个用于 GPT 交互式文本生成的脚本。

训练

数据预处理

训练数据需要预处理。首先，将您的训练数据以松散的 json 格式放置，一个 json 每行包含一个文本样本。例如：

{"src": "www.nvidia.com", "text": "The quick brown fox", "type": "Eng", "id": "0", "title": "First Part"}
{"src": "The Internet", "text": "jumps over the lazy dog", "type": "Eng", "id": "42", "title": "Second Part"}

json的字段名称text可以通过使用中的--json-key标志来更改preprocess_data.py。其他元数据是可选的，在训练中不使用。

然后将松散的 json 处理成二进制格式进行训练。要将 json 转换为 mmap、缓存索引文件或惰性加载器格式，请使用preprocess_data.py. 将--dataset-impl标志分别设置为mmap、cached或lazy（默认为mmap）。为 BERT 训练准备数据的示例脚本是：

python tools/preprocess_data.py \
       --input my-corpus.json \
       --output-prefix my-bert \
       --vocab bert-vocab.txt \
       --dataset-impl mmap \
       --tokenizer-type BertWordPieceLowerCase \
       --split-sentences

输出将是两个名为的文件，在这种情况下，my-bert_text_sentence.bin和my-bert_text_sentence.idx. --data-path在后面的 BERT 训练中指定的是完整路径和新文件名，但没有文件扩展名。

GPT 数据预处理需要进行一些小的修改，即添加合并表、文档结尾标记、删除句子拆分以及更改标记器类型：

python tools/preprocess_data.py \
       --input my-corpus.json \
       --output-prefix my-gpt2 \
       --vocab gpt2-vocab.json \
       --dataset-impl mmap \
       --tokenizer-type GPT2BPETokenizer \
       --merge-file gpt2-merges.txt \
       --append-eod

这里的输出文件被命名为my-gpt2_text_document.bin和my-gpt2_text_document.idx. 和以前一样，在 GPT 训练中，使用不带扩展名的较长名称 as --data-path。

源文件中描述了更多命令行参数preprocess_data.py。

BERT 预训练

bash examples/pretrain_bert.sh

此脚本运行单 GPU 345M 参数 BERT 预训练。调试是单 GPU 训练的主要用途，因为代码库和命令行参数针对高度分布式训练进行了优化。大多数论点都是不言自明的。默认情况下，学习率在训练迭代中线性衰减，从过度迭代--lr设置的最小值开始。用于预热的训练迭代分数由 设置。虽然这是单GPU 训练，但由--min-lr--lr-decay-iters--lr-warmup-fraction--micro-batch-sizeglobal-batch-sizewhcih 是每次迭代的批量大小。数据以 949:50:1 的比例划分为训练/验证/测试集（默认为 969:30:1）。这种分区是即时发生的，但在使用相同随机种子（默认为 1234，或使用 手动指定）的运行中是一致的--seed。我们使用train-iters请求的训练迭代。或者，可以提供--train-samples要训练的样本总数。如果存在此选项，则不需要提供--lr-decay-iters，而是需要提供--lr-decay-samples。

指定了日志记录、检查点保存和评估间隔。检查点激活有助于训练更大的模型和/或批次。请注意，--data-path现在包括_text_sentence在预处理中添加的附加后缀，但不包括文件扩展名。

CHECKPOINT_PATH=checkpoints/bert_345m
VOCAB_FILE=bert-vocab.txt
DATA_PATH=my-bert_text_sentence

BERT_ARGS="--num-layers 24 \
           --hidden-size 1024 \
           --num-attention-heads 16 \
           --seq-length 512 \
           --max-position-embeddings 512 \
           --lr 0.0001 \
           --lr-decay-iters 990000 \
           --train-iters 2000000 \
           --min-lr 0.00001 \
           --lr-warmup-fraction 0.01 \
	   --micro-batch-size 4 \	   
           --global-batch-size 8 \
           --vocab-file $VOCAB_FILE \
           --split 949,50,1 \
           --fp16"

OUTPUT_ARGS="--log-interval 10 \
             --save-interval 500 \
             --eval-interval 100 \
             --eval-iters 10 \
             --checkpoint-activations"

python pretrain_bert.py \
       $BERT_ARGS \
       $OUTPUT_ARGS \
       --save $CHECKPOINT_PATH \
       --load $CHECKPOINT_PATH \
       --data-path $DATA_PATH

源文件中描述了更多命令行参数arguments.py。

GPT 预训练

bash examples/pretrain_gpt.sh

此脚本运行单 GPU 345M 参数 GPT 预训练。如上所述，单 GPU 训练主要用于调试目的，因为代码针对分布式训练进行了优化。

它遵循与之前的 BERT 脚本基本相同的格式，但有一些显着差异：使用的标记化方案是 BPE（需要合并表和json词汇文件）而不是 WordPiece，模型架构允许更长的序列（请注意，最大位置嵌入必须大于或等于最大序列长度），并且--lr-decay-style已设置为余弦衰减。请注意，--data-path现在包括_text_document在预处理中添加的附加后缀，但不包括文件扩展名。

CHECKPOINT_PATH=checkpoints/gpt2_345m
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
DATA_PATH=my-gpt2_text_document

GPT_ARGS="--num-layers 24 \
          --hidden-size 1024 \
          --num-attention-heads 16 \
          --seq-length 1024 \
          --max-position-embeddings 1024 \
          --micro-batch-size 4 \
          --global-batch-size 8 \
          --lr 0.00015 \
          --train-iters 500000 \
          --lr-decay-iters 320000 \
          --lr-decay-style cosine \
          --vocab-file $VOCAB_FILE \
          --merge-file $MERGE_FILE \
          --lr-warmup-fraction .01 \
          --fp16"

OUTPUT_ARGS=<same as those in BERT pretraining above>

python pretrain_gpt.py \
       $GPT_ARGS \
       $OUTPUT_ARGS \
       --save $CHECKPOINT_PATH \
       --load $CHECKPOINT_PATH \
       --data-path $DATA_PATH \

源文件中描述了更多命令行参数arguments.py。

分布式 BERT 或 GPT 预训练

bash examples/pretrain_bert_distributed.sh

bash examples/pretrain_gpt_distributed.sh

这些脚本使用 PyTorch 分布式启动器进行分布式训练。因此，可以通过正确设置环境变量并init_method='env://'在启动器中使用来实现多节点训练。有关这些环境变量的进一步描述，请参阅官方 PyTorch文档。默认情况下，多节点训练使用nccl分布式后端。一组简单的附加参数和使用带有 Python 标志的 PyTorch 分布式模块（详见下文）是采用分布式训练的唯一附加要求。-m torch.distributed.launch

我们使用两种类型的并行：数据并行和模型并行。我们促进了两种分布式数据并行实现：一种是我们自己的简单实现，它在反向传播步骤结束时执行梯度全归约，另一种是 Torch 的分布式数据并行包装器，它将梯度归约与反向传播计算重叠。要在这两个选项之间切换，请使用--DDP-impl local或--DDP-impl torch， 分别。正如预期的那样，Torch 分布式数据并行性在更大的模型尺寸下效率更高。例如，对于在 512 个 GPU 上运行的 83 亿参数模型，当使用 Torch 的分布式数据并行时，缩放比例从 60% 增加到 76%。但是，重叠方法需要更多内存，并且对于某些配置（例如，使用 2 路模型并行的 25 亿个参数和没有模型并行的 12 亿个参数）可能会导致整体训练变慢。我们凭经验发现，在这些情况下使用较小的模型可以缩短训练时间。

其次，我们开发了一种简单有效的二维模型并行方法。要使用张量模型并行性（在多个 GPU 上拆分单个转换器模块的执行），请添加--tensor-model-parallel-size标志以指定要在其中拆分模型的 GPU 数量，以及如上所述传递给分布式启动器的参数。要使用流水线模型并行性（将转换器模块分成多个阶段，每个阶段具有相同数量的转换器模块，然后通过将批处理分解为更小的微批处理来流水线执行），请使用--pipeline-model-parallel-size标志指定将模型拆分为的阶段数（例如，将具有 24 个变换器层的模型拆分为 4 个阶段意味着每个阶段都有 6 个变换器层）。

我们有如何在这些脚本中使用这两种不同形式的模型并行性的示例：

bash examples/pretrain_bert_distributed_with_mp.sh

bash examples/pretrain_gpt_distributed_with_mp.sh

除了这些细微的变化之外，分布式训练与单个 GPU 上的训练相同。

分布式 BERT 训练：

WORLD_SIZE=8
MP_SIZE=2

DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
                  --nnodes 1 \
                  --node_rank 0 \
                  --master_addr localhost \
                  --master_port 6000"

CHECKPOINT_PATH=checkpoints/bert_345m
VOCAB_FILE=bert-vocab.txt
DATA_PATH=my-bert_text_sentence
BERT_ARGS=<same as those in BERT pretraining above>
OUTPUT_ARGS=<same as those in BERT pretraining above>

python -m torch.distributed.launch $DISTRIBUTED_ARGS ./pretrain_bert.py \
                $BERT_ARGS \
                $OUTPUT_ARGS \
                --save $CHECKPOINT_PATH \
                --load $CHECKPOINT_PATH \
                --data-path $DATA_PATH \
                --tensor-model-parallel-size $MP_SIZE \
                --DDP-impl torch

分布式 GPT 训练：

WORLD_SIZE=8
MP_SIZE=2

DISTRIBUTED_ARGS=<same as those directly above>

CHECKPOINT_PATH=checkpoints/gpt2_345m
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
DATA_PATH=my-gpt2_text_document
GPT_ARGS=<same as those in GPT pretraining above>
OUTPUT_ARGS=<same as those in BERT pretraining above>

python -m torch.distributed.launch $DISTRIBUTED_ARGS ./pretrain_gpt.py \
                $GPT_ARGS \
                $OUTPUT_ARGS \
                --save $CHECKPOINT_PATH \
                --load $CHECKPOINT_PATH \
                --data-path $DATA_PATH \
                --tensor-model-parallel-size $MP_SIZE \
                --DDP-impl torch

GPT-3 示例

bash examples/pretrain_gpt3_175B.sh

我们提供了一个示例，说明如何配置 Megatron在 1024 个 GPU 上运行具有 1750 亿个参数的GPT-3 。该脚本是为带有pyxis插件的slurm设计的，但可以很容易地被任何其他调度程序采用。它分别使用 8 路和 16 路张量和管道并行性。使用选项和，训练将从全局批量大小 16 开始，并通过增量步骤 16 将全局批量大小线性增加到 1536，超过 5,859,375 个样本。训练数据集可以是单个集合，也可以是结合一组权重的多个数据集。global-batch-size 1536rampup-batch-size 16 16 5859375

在 1024 个 A100 GPU 上，全局批量大小为 1536，每次迭代大约需要 32 秒，导致每个 GPU 138 teraFLOPs，是理论峰值 FLOPs 的 44%。

评估和任务

我们提供了几个命令行参数，在下面列出的脚本中进行了详细说明，以处理各种零样本和微调的下游任务。但是，您也可以根据需要从其他语料库上的预训练检查点微调您的模型。为此，只需添加--finetune标志并调整原始训练脚本中的输入文件和训练参数。迭代计数将重置为零，优化器和内部状态将重新初始化。如果微调因任何原因中断，请务必--finetune在继续之前移除标志，否则训练将从头开始。

下面针对 GPT 和 BERT 模型描述了几个下游任务。它们可以在分布式和模型并行模式下运行，与训练脚本中使用的更改相同。

GPT 文本生成

bash examples/generate_text.sh

我们主要使用 GPT 预训练脚本生成文本样本。需要做的改动很少，比如我们需要提供预训练检查点的路径、输出样本的长度、是无条件生成文本（--num-samples表示要生成多少样本）还是有条件的（需要通过--sample-input-file <filename>每行该文件将用作条件文本）。可以播放的可选参数很少，例如top-k, top-p, 或greedy(set top-k and top-p to 0) sampling..

CHECKPOINT_PATH=checkpoints/gpt2_345m
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
GPT_ARGS=<same as those in GPT pretraining above>

MAX_OUTPUT_SEQUENCE_LENGTH=1024
TEMPERATURE=1.0
TOP_P=0.9
NUMBER_OF_SAMPLES=2
OUTPUT_FILE=samples.json

python tools/generate_samples_gpt.py \
       $GPT_ARGS \
       --load $CHECKPOINT_PATH \
       --out-seq-length $MAX_OUTPUT_SEQUENCE_LENGTH \
       --temperature $TEMPERATURE \
       --genfile $OUTPUT_FILE \
       --num-samples $NUMBER_OF_SAMPLES \
       --top_p $TOP_P \
       --recompute

GPT 评估

我们在 WikiText 困惑评估和 LAMBADA Cloze 准确性上包含 GPT 评估的示例脚本。

WikiText 困惑度评估

为了与以前的工作进行比较，我们在单词级WikiText-103 测试数据集上评估困惑度，并在使用我们的子词标记器时根据标记的变化适当地计算困惑度。

我们使用以下命令在 345M 参数模型上运行 WikiText-103 评估。

TASK="WIKITEXT103"

VALID_DATA=<wikitext path>.txt
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
CHECKPOINT_PATH=checkpoints/gpt2_345m

COMMON_TASK_ARGS="--num-layers 24 \
                  --hidden-size 1024 \
                  --num-attention-heads 16 \
                  --seq-length 1024 \
                  --max-position-embeddings 1024 \
                  --fp16 \
                  --vocab-file $VOCAB_FILE"

python tasks/main.py \
       --task $TASK \
       $COMMON_TASK_ARGS \
       --valid-data $VALID_DATA \
       --tokenizer-type GPT2BPETokenizer \
       --merge-file $MERGE_FILE \
       --load $CHECKPOINT_PATH \
       --micro-batch-size 8 \
       --checkpoint-activations \
       --log-interval 10 \
       --no-load-optim \
       --no-load-rng

LAMBADA 完形填空精度

为了计算 LAMBADA 完形填空精度（给定前面的标记预测最后一个标记的准确性），我们使用了 LAMBADA数据集的去标记化处理版本。

我们使用以下命令在 345M 参数模型上运行 LAMBADA 评估。请注意，该--strict-lambada标志应用于要求全字匹配。使其lambada成为文件路径的一部分。

TASK="LAMBADA"

VALID_DATA=<lambada path>.json
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
CHECKPOINT_PATH=checkpoints/gpt2_345m
COMMON_TASK_ARGS=<same as those in WikiText Perplexity Evaluation above>

python tasks/main.py \
       --task $TASK \
       $COMMON_TASK_ARGS \
       --valid-data $VALID_DATA \
       --tokenizer-type GPT2BPETokenizer \
       --strict-lambada \
       --merge-file $MERGE_FILE \
       --load $CHECKPOINT_PATH \
       --micro-batch-size 8 \
       --checkpoint-activations \
       --log-interval 10 \
       --no-load-optim \
       --no-load-rng

源文件中描述了更多命令行参数main.py

BERT 任务评估

种族评估

以下脚本微调 BERT 模型以评估RACE 数据集。TRAIN_DATAandVALID_DATA目录包含 RACE 数据集作为单独的文件.txt。请注意，对于 RACE，批量大小是要评估的 RACE 查询的数量。由于每个 RACE 查询有四个样本，因此通过模型的有效批量大小将是命令行上指定的批量大小的四倍。

TRAIN_DATA="data/RACE/train/middle"
VALID_DATA="data/RACE/dev/middle \
            data/RACE/dev/high"
VOCAB_FILE=bert-vocab.txt
PRETRAINED_CHECKPOINT=checkpoints/bert_345m
CHECKPOINT_PATH=checkpoints/bert_345m_race
COMMON_TASK_ARGS="--num-layers 24 \
                  --hidden-size 1024 \
                  --num-attention-heads 16 \
                  --seq-length 512 \
                  --max-position-embeddings 512 \
                  --fp16 \
                  --vocab-file $VOCAB_FILE"

COMMON_TASK_ARGS_EXT="--train-data $TRAIN_DATA \
                      --valid-data $VALID_DATA \
                      --pretrained-checkpoint $PRETRAINED_CHECKPOINT \
                      --checkpoint-activations \
                      --save-interval 10000 \
                      --save $CHECKPOINT_PATH \
                      --log-interval 100 \
                      --eval-interval 1000 \
                      --eval-iters 10 \
                      --weight-decay 1.0e-1"

python tasks/main.py \
       --task RACE \
       $COMMON_TASK_ARGS \
       $COMMON_TASK_ARGS_EXT \
       --tokenizer-type BertWordPieceLowerCase \
       --epochs 3 \
       --micro-batch-size 4 \
       --lr 1.0e-5 \
       --lr-warmup-fraction 0.06

MNLI 评估

以下脚本微调 BERT 模型以使用MultiNLI 句子对语料库进行评估。由于匹配任务非常相似，因此可以快速调整脚本以与Quora Question Pairs (QQP) 数据集一起使用。

TRAIN_DATA="data/glue_data/MNLI/train.tsv"
VALID_DATA="data/glue_data/MNLI/dev_matched.tsv \
            data/glue_data/MNLI/dev_mismatched.tsv"
PRETRAINED_CHECKPOINT=checkpoints/bert_345m
VOCAB_FILE=bert-vocab.txt
CHECKPOINT_PATH=checkpoints/bert_345m_mnli
COMMON_TASK_ARGS=<same as those in RACE Evaluation above>
COMMON_TASK_ARGS_EXT=<same as those in RACE Evaluation above>

python tasks/main.py \
       --task MNLI \
       $COMMON_TASK_ARGS \
       $COMMON_TASK_ARGS_EXT \
       --tokenizer-type BertWordPieceLowerCase \
       --epochs 5 \
       --micro-batch-size 8 \
       --lr 5.0e-5 \
       --lr-warmup-fraction 0.065

数据集

我们不托管任何用于 GPT 或 BERT 训练的数据集，但是，我们会详细说明它们的集合，以便可以复制我们的结果。

收集维基百科训练数据

我们建议遵循 Google 研究指定的 Wikipedia 数据提取过程：“推荐的预处理是下载最新的转储，使用WikiExtractor.py提取文本，然后应用任何必要的清理将其转换为纯文本。”

我们建议--json在使用 WikiExtractor 时使用该参数，它将 Wikipedia 数据转储为松散的 json 格式（每行一个 json），使其在文件系统上更易于管理，并且也易于被我们的代码库使用。我们建议通过 nltk 标点标准化进一步预处理这个 json 数据集。对于 BERT 训练，使用上述标志--split-sentencesto在生成的索引中包含断句。如果您想将 Wikipedia 数据用于 GPT 训练，您仍应使用 nltk/spacy/ftfy 清理它，但不要使用该标志。preprocess_data.py--split-sentences

收集 GPT Web 文本数据

我们利用来自jcpeterson和eukaryote31 的公开可用的OpenWebText库来下载 url。然后，我们根据我们的openwebtext目录中描述的过程过滤、清理和删除所有下载的内容。对于截至 2018 年 10 月的内容对应的 reddit URL，我们得到了大约 37GB 的内容。