英伟达1小时成功训练BERT，83亿参数打造史上最大语言模型

英伟达一举创造了2个壮举！训练出了世界上最大的语言模型——MegatronLM，包含83亿参数，比BERT大24倍，比GPT-2大5.6倍；还打破了实时对话AI的记录，仅耗时53分钟即可训练出行业标准BERT模型、2毫秒就能对答案做出推断！

世界上最大的语言模型来了，顺便还破了个记录！

英伟达宣布，目前已经训练出了世界上最大的语言模型——MegatronLM。

这个模型有多大？83亿个参数！比谷歌的 BERT 大24倍，比 OpenAI 的 GPT-2 大5.6倍！

不仅如此，英伟达还宣布打破了实时对话 AI 的记录——耗时53分钟就可以训练出行业标准的BERT模型、2毫秒左右就能对答案做出推断。

为了实现这一壮举，英伟达利用模型的并行性，将一个神经网络分割成多个部分，创建了因数据太大无法容纳在单个GPU的训练模型。

最重要的是，代码已开源！

GitHub项目地址：
https://github.com/NVIDIA/Megatron-LM  

MegatronLM，堪称NLP 界的“威震天”！

有钱任性：训练史上最大语言模型需要多少GPU？  

更大的语言模型对于诸如文章完成、问题回答和对话系统等NLP任务非常有用。最近，训练最大的神经语言模型已经成为提高NLP应用水平的最佳方法。

最近的两篇论文，BERT和GPT-2，展示了大规模语言建模的好处。这两篇论文都利用了计算机和可用文本语料库的进步，在自然语言理解、建模和生成方面显著超越了当前的最优水平。

训练这些模型需要数以百计exaflops级的计算力和巧妙的内存管理，以换取减少内存占用的重新计算。然而，对于超过10亿参数的超大型的模型，单个GPU上的内存不足以匹配模型以及训练所需的参数，需要利用模型并行性来将参数分割到多个GPU上。有几种建模并行性的方法，但是它们很难使用，因为它们依赖于自定义编译器，或者扩展性很差，或者需要对优化器进行更改。

在这项工作中，我们通过对现有PyTorch transformer实现进行少量有针对性的修改，实现了一种简单而有效的模型并行方法。我们的代码是用原生Python编写的，利用混合精度训练，并利用NCCL库在GPU之间进行通信。     我们通过在512个GPU上训练一个transformer语言模型证明了这种方法的有效性，该模型具有8路模型并行性和64路数据并行性，83亿参数，使其成为有史以来规模最大的基于transformer的语言模型，其大小为BERT的24倍，GPT-2的5.6倍。我们已经在GitHub存储库中发布了实现此方法的代码。

我们的实验是在英伟达的DGX SuperPOD上进行的。在没有模型并行性的情况下，我们可以在单个V100 32GB GPU上训练一个12亿参数的基线模型，并在整个训练过程中保持39 TeraFLOPS，这是DGX2-H服务器上单个GPU理论峰值的30%。

我们将模型参数扩展到83亿，使用512个GPU，通过8路模型并行化，在整个应用程序中我们实现了高达15.1 PetaFLOPS的持续性能，与单GPU相比，扩展效率达到76%。图1显示了扩展的结果。    

图1：模型并行(蓝色)：多达8路模型并行弱扩展，每个GPU大约有10亿个参数(例如2个GPU有20亿参数，4个GPU有40亿参数)。模型+数据并行(绿色)：类似于模型并行的64路数据并行的配置。

多GPU并行性

训练模型的典型范例是利用 weak scaling 方法和分布式数据并行性，根据GPU的数量来扩展训练批大小。这种方法允许模型在更大的数据集上进行训练，但有一个约束，即所有参数必须适合一个GPU。

模型并行训练可以通过跨多个GPU划分模型来克服这一限制。近年来出现了几个通用模型并行框架，如GPipe和Mesh-TensorFlow。gPipe在不同的处理器上划分层组，而Mesh-TensorFlow使用层内模型并行性。我们的方法在概念上类似于Mesh-TensorFlow，我们关注层内并行性并融合GEMM以减少同步。然而，我们只对现有PyTorch transformer实现进行了一些有针对性的修改，以便使用模型并行性来训练大型transformers。我们的方法很简单，不需要任何新的编译器或代码重新连接来实现模型并行性，并且可以通过插入一些简单的primitives(图2中的f和g 算子)完全实现。

我们利用 transformer网络的结构，通过添加一些同步primitives来创建一个简单的模型并行实现。

transformer层由一个self attention block和一个2层的多层感知器(MLP)组成。我们分别在这两个模块中引入模型并行性。

如图2a所示，这是MLP的结构，由两个GEMM组成，中间有一个GeLU非线性，后面有一个dropout层。我们以列并行方式划分第一个GEMM。这使得GeLU 非线性可以独立地应用于每个分块GEMM的输出。模块中的第二个GEMM沿着行并行化，直接获取GeLU层的输出，不需要任何通信。然后，在将输出传递到dropout层之前，跨GPU减少第二个GEMM的输出。这种方法将MLP block中的GEMM跨GPU分割了，只需要在正向传递(g算子)中执行一个all-reduce操作，在反向传递(f算子)中执行一个all-reduce操作。    

图2：(a): MLP， (b)：transformer的self attention block。

如图2(b)所示，在self attention block上，我们利用multihead attention操作中的固有并行性，以列并行方式划分与键（K），查询（Q）和值（V）相关联的 GEMM。

这使得我们可以在GPU之间分割每个attention head参数和工作负载，并且不需要任何即时通信来完成self attention。

这种方法对于MLP和self-attention层都融合了两个GEMM的组，消除了中间的同步点，并获得了更好的scaling性能。这使我们能够在一个简单的transformer层中执行所有GEMM，只使用前向路径的2个all reduce和后向路径的2个all reduce，如图3所示。

图3：GPT-2 transformer层的模型并行性。

这种方法实现起来很简单，因为它只需要在向前和向后传递中添加一些额外的all-reduce操作。它不需要编译器，并且与gPipe等方法提倡的那种pipeline模型并行性是正交的。

性能

为了测试我们的实现的计算性能，我们考虑了表1中四组参数的GPT-2模型。

表1：用于scaling 研究的参数。

所有的实验都是在NVIDIA的DGX SuperPOD上进行的，我们使用了多达32台DGX- 2h服务器(总共512个Tesla V100 SXM3 32GB GPU)。该系统针对多节点深度学习应用程序进行了优化，服务器内部GPU之间的带宽为300 GB/s，服务器之间的互连带宽为100 GB/s。

图4显示了模型和模型+数据并行性的扩展值。我们在这两种设置中都观察到了出色的扩展数字。例如，8路(8 GPU)模型并行的83亿参数模型实现了77%的线性扩展。模型+数据并行性要求在反向传播步骤之后进一步通信梯度，因此扩展数略有下降。然而，即使是运行在512个GPU上的最大配置(83亿参数)，相对于强大的基准单GPU配置(12亿个参数)，我们仍然可以实现74%的扩展性。

图4：模型(左)和模型+数据(右)随着GPU的数量并行地进行weak scaling。

最后，我们研究了attention heads对模型并行扩展的影响。为此，我们考虑了83亿参数、具有8路模型并行性的参数配置，并将attention heads的数目从16个改为32个。结果如表2所示。随着attention heads数量的增加，self attention层中的一些GEMM变小，同时softmax中的元素数量增加。这导致了轻微的scaling decrease。未来的研究在设计大型transformer模型时应该警惕这种超参数，平衡模型性能和模型效率。  

表2：attention heads 数量对scaling的影响。  

GPT-2训练

为了训练GPT-2模型，我们创建了一个从_Reddit_下载的37 GB _WebText_ dataset，它类似于原始GPT-2论文中描述的webtext数据集。数据集最终有810万个url。我们将WebText数据集随机分割为95:5的比例，分别得到训练集和验证集。我们考虑了4种参数规模的模型：3.45亿、7.75亿、25亿和83亿。

图5：训练子集的验证困惑度。在对37GB数据集过拟合之后，8.3B模型提前停止了。

图5显示了验证的困惑度(perplexity)。我们发现。最大的83亿参数的语言模型在~6epoch之后开始overfit，一种1 epoch被定义为15200次迭代。我们认为这可以通过使用更大规模的数据集来缓解，类似于XLNet和RoBERTa等最近论文中使用的数据集。

GPT-2评估

为了分析大型语言模型的训练性能，我们在wikitext-103数据集上计算了perplexity，在Lambada数据集上计算了closize风格的预测精度。

正如预期的一样，wikitext perplexity随着模型尺寸的增大而减小，lambada准确率随着模型尺寸的增大而增加(表3)。

表3：wikitext perplexity(越低越好)和Lambada完形精度(越高越好)的评估结果。

结论

在这项工作中，我们在现有的深度学习硬件、软件和模型的基础上，构建了世界上最大的基于transformer的语言模型。

在此过程中，我们成功地突破了传统的单GPU训练的限制，实现了一种简单而高效的模型并行方法，只需对现有PyTorch transformer实现进行少量有针对性的修改。

我们在512台NVIDIA V100 GPU上高效地训练了83亿参数的语言模型(分别比BERT和GPT-2大24倍和5.6倍)，具有8路模型并行性，并在整个应用程序中实现了高达15.1千万亿次浮点运算(PetaFLOPS)。

我们发现，与较小的transformer模型相比，更大的transformer模型可以在相同的时间内进行训练，并且可以显著提高性能。

然而，正如我们在工作中所展示的，NLP仍然需要合适的数据集、问题和技术来正确地训练这些大型语言模型，否则会出现过拟合。

我们将我们的工作开源，以便社区就可以复制并扩展它们。

英伟达官方GitHub项目已开源！

英伟达在官方GitHub上对MegatronLM开源了代码，也提供了相应的教程。

项目地址：https://github.com/NVIDIA/Megatron-LM

安装

官方只支持 Python 3.6。请安装支持GPU的最新版本PyTorch。

此外，代码库的一部分利用tensorflow-cpu（可选）执行TFRecords的数据加载以进行BERT训练。

建议要么使用./docker/中提供的Dockerfile，要么创建一个虚拟环境(以避免破坏现有的tf安装)并安装requirements.txt。

1python -m pip install virtualenv 2virtualenv bert_env 3source bert_env/bin/activate 4pip install -r requirements.txt

用法

提供了5个预训练BERT的脚本和3个预训练GPT2的脚本。使用 --save 和 --load 保存并加载模型检查点(checkpoint)。

此外，还提供 GPT2 脚本，用于在wiki文本和LAMBADA上生成GPT2的交互式文本生成和零样本(zero shot)评估。

BERT预训练

1bash scripts/pretrain_bert.sh

此脚本运行单个gpu BERT预训练，主要用于调试目的。优化参数设置为64路分布式训练。

要使用此脚本，请 --train-data以loose json格式放置，每行一个json。json字典的文本字段应该对应于 --text-key。

1python pretrain_bert.py  2       --num-layers 24  3       --hidden-size 1024  4       --num-attention-heads 16  5       --batch-size 4  6       --seq-length 512  7       --max-preds-per-seq 80  8       --max-position-embeddings 512  9       --train-iters 1000000  10       --save checkpoints/bert_345m  11       --load checkpoints/bert_345m  12       --resume-dataloader  13       --train-data wikipedia  14       --lazy-loader  15       --tokenizer-type BertWordPieceTokenizer  16       --tokenizer-model-type bert-large-uncased  17       --presplit-sentences  18       --cache-dir cache  19       --split 949,50,1  20       --distributed-backend nccl  21       --lr 0.0001  22       --lr-decay-style linear  23       --lr-decay-iters 990000  24       --weight-decay 1e-2  25       --clip-grad 1.0  26       --warmup .01  27       --fp16  28       --fp32-embedding

GPT2 预训练

1bash scripts/pretrain_gpt2.sh

此脚本运行单gpu gpt2预训练，主要用于调试目的。优化参数设置为64路分布式训练。

它与前一个脚本格式大致相同，但有一些值得注意的差异：

--tokenizer-type已切换为GPT2BPETokenizer；

--lr-decay-style已切换为cosine decay等等。

另外，GPT2使用来自BERT的不同参数初始化，用于训练深度残差网络。要使用此初始化来训练BERT，请使用--deep-init。

1python pretrain_gpt2.py  2       --num-layers 24  3       --hidden-size 1024  4       --num-attention-heads 16  5       --batch-size 8  6       --seq-length 1024  7       --max-position-embeddings 1024  8       --train-iters 320000  9       --save checkpoints/gpt2_345m  10       --load checkpoints/gpt2_345m  11       --resume-dataloader  12       --train-data wikipedia  13       --lazy-loader  14       --tokenizer-type GPT2BPETokenizer  15       --cache-dir cache  16       --split 949,50,1  17       --distributed-backend nccl  18       --lr 0.00015  19       --lr-decay-style cosine  20       --weight-decay 1e-2  21       --clip-grad 1.0  22       --warmup .01  23       --checkpoint-activations  24       --fp16