【新智元导读】 英伟达一举创造了 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 毫秒左右就能对答案做出推断。
对话AI的下一步是什么？https://www.zhihu.com/video/1144966104484024320
为了实现这一壮举，英伟达利用模型的并行性，将一个神经网络分割成多个部分，创建了因数据太大无法容纳在单个 GPU 的训练模型。
最重要的是，代码已开源！
GitHub 项目地址：
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 开源了代码，也提供了相应的教程。
安装
官方只支持 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
更多细节内容，读者可前往官方 GitHub 浏览：github.com/NVIDIA/Mega…
