搞定千亿参数，训练时间只用1/3，微软全新工具催生超级NLP模型

GPU 显存一直是训练先进深度学习模型的最大瓶颈——大规模训练经常会碰到模型参数太多，显存无法容纳的情况。

最近，微软提出的全新方法 DeepSpeed 将每个节点的批处理大小增大了 4 倍，而训练时间减少到原来的 1/3，新的工具兼容深度学习框架 PyTorch。这一研究会成为深度学习发展的重要一步吗？

DeepSpeed 的强大性能也催生了全新的语言生成模型：微软同时发布了一款参数量高达 170 亿的模型 T-NLG，它在多项任务中打破了以往由 GPT-2 等保持的业内最佳记录。

Turing-NLG：170 亿参数，史上最大的 NLG 模型

在现有的 NLP 模型中，Turing Natural Language Generation（T-NLG）的参数量远超其他。

诸如 BERT 和 GPT-2 这样的深度学习语言模型已经在大部分 NLP 下游任务中实现了 SOTA，包括问答、对话代理和文档理解等。表现良好的自然语言生成（NLG）模型可以部署到多种应用程序，比如帮助作者进行内容撰写，可节省用户时间的长文本总结功能，或者使用智能助手改善用户体验等等。

近日，微软在最新的博客文章中介绍了史上最大的 NLG 模型「T-NLG」，参数量达到 170 亿。在各项语言模型基准上，T-NLG 均实现了 SOTA 表现，同时在问答和总结这样的实际任务中也表现出色。

模型架构

T-NLG 是基于 Transformer 的语言生成模型，这意味着它能够生成单词来完成开放式文本任务。除了补全不完整的句子以外，它还可以针对问题生成直接答案以及针对输入文档生成摘要等。

模型越大，预训练数据也会更加多样及复杂，即使训练实例比较少，它也可以在下游任务的扩展中拥有良好表现。因此，训练一个大型集中式多任务模型并在多种任务中共享功能，会比针对每个任务单独训练新模型要更高效一些。

训练方法

超过 13 亿参数的模型无法部署在单个 GPU 上，因此该模型需要在多个 GPU 之间并行化或者被分解。训练 T-NLG 的过程有以下几项软件和硬件上的突破：

1. 将 NVIDIA DGX-2 与 InfiniBand 连接起来，使得 GPU 之间的通信更快；

2. 使用张量切片（tensor slicing）在 NVIDIA Megatron-LM 框架上的四个 NVIDIA V100 GPU 上进行模型分割；

3. 带有 ZeRO 优化器的 DeepSpeed 使得模型并行度从 16 降至 4，并将每个节点的批处理大小增至 4 倍，同时训练时间减至三分之一。DeepSpeed 可以在更少 GPU 的前提下，做到让大型模型的训练效率更高，并且仅使用 256 个 NVIDIA GPU 就可以进行 512 的批量训练，相比之下，单独使用 Megatron-LM 需要 1024 个 NVIDIA GPU。DeepSpeed 与 PyTorch 兼容。

生成的 T-NLG 模型有 78 个 Transformer 层，包括 4256 个隐藏层参数以及 28 个注意力头。为使生成结果更接近于 Megatron-LM，团队使用了与 Megatron-LM 相同的超参数和 Learning schedule 对模型进行训练，使用了 1024 长度的 Token 序列，训练了 30 万步，批大小为 512，损失函数为自回归生成损失（autoregressive generation loss）。

Learning schedule 遵循 3200 个线性 warmup 步骤，最高学习速率为 1.5x10^-4，余弦衰减超过 50 万步，16 位浮点数，训练数据类型与 Megatron-LM 相同。

训练数据

同时，团队也在一系列标准语言任务中对比了 T-NLG 预训练模型，比如像 WikiText-103 的困惑度（越低越好），还有 LAMBADA 文本词汇预测准确性（越高越好）等任务。

如下表所示，T-NLG 在上述两项任务中实现了 SOTA，Megatron-LM 的数据来源于 NVIDIA Megatron 模型的公开结果。

此处 Open AI 使用了额外的停用词过滤处理程序，实现了比单独实现模型更高的数字。Megatron-LM 和 T-NLG 则没有使用。

下图展示了 T-NLG 与 Megatron-LM 在困惑度验证方面的对比表现。

二者训练过程中的困惑度验证表现对比：橙色线代表 Megatron-8B 参数模型，蓝/绿色线代表 T-NLG。虚线代表的是现有 SOTA 模型能实现的最低验证损失。从蓝色到绿色的过渡意味着 T-NLG 在性能上超越了现有的其他训练模型。

实验结果

机器问答

在搜索引擎上提出一个问题的时候，很多用户习惯于在结果页面顶端看到答案卡片。多数卡片会展示一个含有答案的句子，连带给出这个句子的原始语境和出处，如下所示：

相比之下，T-NLG 会直接用一个完整的句子来回答问题（不带原始语境），这点对于 AI 助手非常重要，因为它可以帮助它们更加智能地回答用户有关自己个人数据、邮件、Word 文档等方面的问题。

此外，T-NLG 还能完成「zero shot」式的问答，即在没有语境段落的情况下给出答案。在下面这个例子中，研究者就只给出了问题，没有给出语境。在这些情况下，模型依赖于预训练得到的知识来生成答案。

此外，研究人员还请人类标注员来评估 T-NLG 模型与之前的基线模型，结果如下：

研究人员还注意到，虽然模型变大了，但它学习所需的下游任务实例反而更少。他们只有不超过 10 万组直接能够表示为「问题-段落--答案」的三元组。但即使只经过了几千个实例的训练，他们得到的模型也优于 LSTM 基线，后者是在相同数据的多个 epoch 上训练的。这一结果有很高的商业价值，因为收集带标注的监督数据成本非常高。

文本摘要

T-NLG 的目标是写出接近人类的文本摘要（email、博客、Word 文档、Excel 图表等），但监督训练数据的缺乏阻碍了这一目标的实现。T-NLG 的优势在于，它已经非常善于理解文本，因此不需要太多的监督数据就能超越之前使用的技术。

为了使 T-NLG 适用于更多的文本类型，研究人员在几乎所有公开可用的公共摘要数据集上微调了 T-NLG 模型，使其具备适应多任务的能力。训练所用的示例接近 400 万。

下图展示了 T-NLG 与之前的 SOTA 模型以及 PEGASUS（最近的一个基于 Transformer 的模型）的 Rouge 分数对比：

下图展示了 T-NLG 任务在一些文本摘要任务中的应用示例：

ZeRO＆DeepSpeed 助力

T-NLG 模型之所以能够实现快速的训练和收敛，需要归功于团队在模型加速方面的研究。在这一方面，微软已推出了一个名为 DeepSpeed 的代码库，这一代码库能够和 PyTorch 兼容。同时和 DeepSpeed 代码库一起推出的还有名为 ZeRO 的优化方法。

同时，该库中有一个很重要的组件：ZeRO。这是一个新的并行优化组件，能够很好地降低模型和数据并行需要的资源，并提升可训练的参数量。研究者将这个组件用在了 T-NLG 模型的训练上。

据悉，ZeRO 对于内存的优化可以让 1000 亿参数的模型在现有的 GPU 集群上提升 3 到 5 倍的吞吐量。

现有的加速方法的挑战

为什么要推出新的模型加速库呢？微软认为，现有的解决方法在计算、通信和开发效率方面牺牲了很多。首先，数据并行并不能降低每个设备的内存占用情况，一个超过 10 亿参数的模型占用的内存量甚至会超出 32G 内存的 GPU。

其次，由于高细粒度的计算和高昂的通信成本，模型并行并不能扩展超过一个单节点的效率。例如，英伟达的 Megatron-LM 实现了 83 亿参数的模型加速。在单节点多 GPU 上，这个模型可以很好地适应，但是要扩展到多个节点的时候，模型性能就下降了。

ZeRO 的解决方法

为了解决数据并行和模型并行存在的问题，ZeRO 提供了三阶段的优化方法，分别为优化器状态分割、梯度分割、参数分割，三个阶段按顺序实施。

1. 在优化器分割状态：ZeRO 降低了 3/4 的内存，通信量和数据并行相同；

2. 加入梯度分割：降低了 7/8 的内存，通信量和数据并行相同；

3. 加入参数分割：内存减少与数据并行度呈线性关系。例如，在 64 个 GPU 上进行分割的时候，可以将内存降至 1/64。在通信量上有 50% 的提升。

据悉，ZeRO 可以训练兆级别参数量的模型，使用 1024 个英伟达 GPU。对于一个拥有一兆参数、使用 Adam 优化器的模型来说，如果精度要求是 16 比特，那么它需要大约 16TB 的内存提供给优化器状态、梯度和参数。16T 除以 1024 为 16G，对于一个 GPU 而言就更为合理了。

图 1：与标准数据并行基准相比，ZeRO 在三个阶段中节省的内存和通信用量。

ZeRO 和 DeepSpeed 的结合

ZeRO 和 DeepSpeed 结合的优势：

• 规模：OpenAI 的 GPT-2、英伟达的 Megatron-LM 和谷歌的 T5（当前 SOTA 模型）参数量分别为 15 亿、83 亿和 110 亿。ZeRO-OS 提供的系统支持可以运行 1000 亿参数的模型；

• 速度：在不同硬件上，吞吐是当前 SOTA 的 5 倍；

• 成本：吞吐量提升能够大幅度降低训练成本；

• 易用性：只需更改几行代码就能让 PyTorch 模型使用 DeepSpeed 和 ZeRO 成为可能。与当前模型并行库相比，DeepSpeed 库不需要重新设计代码或重构模型。同时，DeepSpeed 库也不对模型尺寸（如注意力头的数量、隐藏层参数量等）、批大小或任何其他训练参数施加任何限制。

DeepSpeed 开源库：

兼容 PyTorch，实现系统性能

研究者实现了第一阶段的 ZeRO，即优化器状态分区（optimizer states partitioning，ZeRO-OS），它能够支持 1000 亿参数的模型。相关代码将与 DeepSpeed 训练优化库一同发布。借助于兼容 PyTorch 的轻量级 API，DeepSpeed 库提供了当前 SOTA 的训练方法，如 ZeRO、分布式训练、混合精度和检查点技术等。

因此，只需更改 PyTorch 模型的几行代码，用户就可以利用 DeepSpeed 库来解决底层性能问题，还能够提升训练速度和规模。

DeepSpeed+ZeRO 在 T-NLG 上的应用

研究者在 DeepSpeed 开源库中联合使用 ZeRO-OS 和英伟达的 Megatron-LM 来训练 T-NLG 模型。与单独使用英伟达 Megatron-LM 相比，二者结合所节省的内存能够使得 Turning-NLG 模型的并行度降低 3/4，批大小增加到 4 倍，因而可以实现 3 倍的吞吐量增益。

此外，在批大小为 512 时，ZeRO-OS 和 Megatron-LM 的结合可以将所需的 GPU 数量减少到 256 个，而单独使用 Megatron-LM 则需要 1024 个 GPU。最后，Megatron-LM 不能单独运行标准的 T-NLG，因为它的注意力头（=28）不能被模型并行度（=16）整除，所以不支持该模型结构。

潜在应用

T-NLG 现在已经实现了在自然语言生成领域的 SOTA，这对于微软的很多业务和客户都是有利的。例如，模型可以帮助用户对文档和邮件进行摘要化处理。同时，模型也能够应用在 Office 系列软件中，比如作为文档的写作助理或问答机器人，分别面向写作者和读者。另外，微软也十分看好这一模型在客服机器人等方面的潜力。

在未来，微软会继续探索深度学习方面的应用，并将研究成果和旗下的 Bing、Office、Xbox 等方面的软硬件服务结合起来。

原文链接：

1. https://www.microsoft.com/en-us/research/blog/turing-nlg-a-17-billion-parameter-language-model-by-microsoft/

2. https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/

本文由机器之心编译自微软研究院博客。

你也许还想看：