世界上最大的语言模型来了，顺便还破了个记录！
英伟达宣布，目前已经训练出了世界上最大的语言模型--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。
1 --python 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
更多细节内容，读者可前往官方GitHub浏览：
https://github.com/NVIDIA/Megatron-LM
参考链接：
VB：https://venturebeat.com/2019/08/13/nvidia-trains-worlds-largest-transformer-based-language-model/
TechCrunch：https://techcrunch.com/2019/08/13/nvidia-breaks-records-in-training-and-inference-for-real-time-conversational-ai/
GitHub：https://github.com/NVIDIA/Megatron-LM
