英伟达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-mpipinstallvirtualenv 2virtualenvbert_env 3sourcebert_env/bin/activate 4pipinstall-rrequirements.txt

用法

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

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

BERT预训练

1bashscripts/pretrain_bert.sh

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

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

1pythonpretrain_bert.py 2--num-layers24 3--hidden-size1024 4--num-attention-heads16 5--batch-size4 6--seq-length512 7--max-preds-per-seq80 8--max-position-embeddings512 9--train-iters1000000 10--savecheckpoints/bert_345m 11--loadcheckpoints/bert_345m 12--resume-dataloader 13--train-datawikipedia 14--lazy-loader 15--tokenizer-typeBertWordPieceTokenizer 16--tokenizer-model-typebert-large-uncased 17--presplit-sentences 18--cache-dircache 19--split949,50,1 20--distributed-backendnccl 21--lr0.0001 22--lr-decay-stylelinear 23--lr-decay-iters990000 24--weight-decay1e-2 25--clip-grad1.0 26--warmup.01 27--fp16 28--fp32-embedding

GPT2 预训练

1bashscripts/pretrain_gpt2.sh

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

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

--tokenizer-type已切换为GPT2BPETokenizer；

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

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

1pythonpretrain_gpt2.py 2--num-layers24 3--hidden-size1024 4--num-attention-heads16 5--batch-size8 6--seq-length1024 7--max-position-embeddings1024 8--train-iters320000 9--savecheckpoints/gpt2_345m 10--loadcheckpoints/gpt2_345m 11--resume-dataloader 12--train-datawikipedia 13--lazy-loader 14--tokenizer-typeGPT2BPETokenizer 15--cache-dircache 16--split949,50,1 17--distributed-backendnccl 18--lr0.00015 19--lr-decay-stylecosine 20--weight-decay1e-2 21--clip-grad1.0 22--warmup.01 23--checkpoint-activations 24--fp16