有钱任性：英伟达训练80亿参数量GPT-2，1475块V100 53分钟训练BERT

anrey573

机器之心报道
机器之心编辑部

你见过 50 分钟训练 BERT-Large、80 亿参数量训练 GPT-2 吗？快看看经过 CUDA 优化的 Transformer 为什么这么强。

英伟达今日宣布，该公司打破了 NLP 领域的三项记录：
1）将 BERT 的训练时间缩短到了 53 分钟；2）将 BERT 的推理时间缩短到了 2.2 毫秒（10 毫秒已经是业界公认的高水平）；3）将 GPT-2 的参数量推向 80 亿（以前 OpenAI GPT-2 最大为 15 亿参数量）。



这些突破可以为现实世界中所有使用 NLP 对话 AI 和 GPU 硬件的用户带来很多便利，如降低语音助手的反应延时，使其与人类的交流更加自然。



训练最快的语言模型
英伟达在 BERT 训练、推理时间上的突破离不开其 SuperPOD 系统。它由 92 个英伟达 DGX-2H 系统组成，运行在 1472 块 V100 GPU 上。该系统在 T4 GPU 上管理推断行为，它的性能甚至比高度优化的 CPU 还要好几个数量级。
英伟达使用 PyTorch 运行整个 BERT-Large 模型，并采用了自动混合精度方法以加速吞吐量。对于一般的研究者，只要有一个 DGX-2 服务器（16 块 V100），我们就能在 3 天内完成 BERT-Large 模型的训练。如下展示了在不同 GPU 数量下的训练时长：



目前，英伟达已经开源了 BERT 的训练代码以及 TensorRT 优化的 BERT 样本，地址和下面的预训练 GPT-2 是一样的。
训练最大的语言模型
如果训练变得更快，那么这意味着什么？当然是我们能训练更大的模型啦。英伟达成功地构建并训练了最大的语言模型 GPT-2 8B，这一模型包含 83 亿参数量，是 BERT-Large 模型的 24 倍、GPT-2 的 5.6 倍。想想我们一个 BERT-Large 都训练不了，英伟达还训练「24 个」BERT-Large，这也是很优秀了。



英伟达将这一模型称为「Megatron」（威震天），还开源了用来训练这一模型的 pytorch 代码。

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

所以这样超大模型的结构应该是什么样的？当然基础结构还是 Transformer，但是超参的配置确实非常惊人，你见过 72 层、每层隐藏单元都是 3072 的 Treansformer 么。。



这么巨大模型，当然效果也还是挺好的，83 亿参数量的 GPT-2 在验证困惑度上，下降地非常快，差不多 5 个 Epoch 就能到达非常理想的效果。如下所示为模型大小与验证困惑度之间的关系。



为了训练如此庞大的模型，英伟达利用了模型并行，用创建大模型的技术将模型分成几个部分。利用这一技术创建的模型非常大，单个 GPU 内存不足以应付，所以需要模型并行来分解压力。
虽然模型并行本来就有一些开销，且如果模型能装进单块 GPU 最好不用模型并行，但是对于 80 亿参数量的 GPT-2，模型并行就是必不可少的。如下展示了随 GPU 数量的增加，其所提供有效计算力的增长差不多接近线性。



为什么这个 Transformer 能这么快
前面不论是训练 BERT-Large 还是巨型 GPT-2，它们的基础都是 Transformer，如果 Transformer 训练得不够快，那么堆再多的算力也不能完成这两项挑战。Faster Transformer 就是支持极速训练极大模型的基础，与两个模型开源的同时，英伟达也开源了 Faster Transformer，读者可以获取项目全部源代码，最新的性能数据以及支持的特性。

开源地址：https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer

如下英伟达 GPU 计算专家团队贾晓莹将向各位读者介绍 Faster Transformer 这一制胜武器。
什么是 Faster Transformer
目前 Transformer 在多种场景下都有非常优秀的表现，但是在推理部署阶段，其计算性能却受到了巨大的挑战：以 BERT 为原型的多层 Transformer 模型，其性能常常难以满足在线业务对于低延迟（保证服务质量）和高吞吐（考虑成本）的要求。以 BERT-BASE 为例，超过 90% 的计算时间消耗在 12 层 Transformer 的前向计算上。
因此，一个高效的 Transformer 前向计算方案，既可以为在线业务带来降本增效的作用，也有利于以 Transformer 结构为核心的各类网络在更多实际工业场景中落地，这也就是 Faster Transformer 诞生的前提。
Faster Transformer 是一个基于 CUDA 和 cuBLAS 的 Transformer Encoder 前向计算实现，其代码简洁明了，后续可以通过简单修改支持多种 Transformer 结构。目前优化集中在编码器的前向计算（解码器开发在后续特性规划中）。底层由 CUDA 和 cuBLAS 实现，支持 FP16 和 FP32 两种计算模式，其中 FP16 可以充分利用 Volta 和 Turing 架构 GPU 上的 Tensor Core 计算单元。 Faster Transformer 共接收 4 个输入参数。首先是 attention head 的数量以及每个 head 的维度，这两个参数是决定 Transformer 网络结构的关键参数。这两个参数的动态传入，可以保证 Faster Transformer 既支持标准的 BERT-BASE（12 head x 64 维），也支持裁剪过的模型（例如，4 head x 32 维），或者其他各式专门定制化的模型。
其余两个参数是 Batch Size 和句子最大长度。出于性能考虑，目前句子最大长度固定为最常用的 32，64 和 128 三种，未来会支持任意长度。Faster Transformer 对外提供 C++ API，TensorFlow OP 接口，以及 TensorRT 插件，并提供了相应的示例，用以支持用户将其集成到不同的线上应用代码中。
Faster Transformer 性能怎么样
Faster Transformer 在不同的应用场景下都有着突出的表现。英伟达在这里测试了不同生产环境下 Faster Transformer 前向计算的执行时间以及与 TensorFlow XLA 的性能比较。
英伟达考察了 Faster Transformer 在搜索或者广告推荐等大 batch size 场景下的加速效果。下面两张表分别测试了固定句子长度为 32，标准模型（12 head x 64 维）和裁剪模型（4 head x 32 维）在不同 batch size 下，12 层 Transformer 在 V100 上使用了 FP16 计算精度的性能。


表：标准模型不同 Batch Size 下 TensorFlow XLA 和 Faster Transformer 在 V100 上的性能对比



表：裁剪模型不同 Batch Size 下 TensorFlow XLA 和 Faster Transformer 在 V100 上的性能对比

可以看出，在标准模型和裁剪模型上，Faster Transformer 都有很好的加速效果。
Faster Transformer 优化原理
Faster Transformer 提供了 TensorFlow OP，C++ API 和 TensorRT Plugin 三种接口。在 TensorFlow 中使用 Faster Transformer 最为简单。只需要先 import .so 文件，然后在代码段中添加对 Faster Transformer OP 的调用即可。
虽然使用起来非常便捷，但 Faster Transformer 的优化原理是什么样的？
在深入了解 Faster Transformer 的优化原理之前，我们先来看下 TensorFlow 的实现情况。下图展示了 TensorFlow 在默认计算模式（不使用 XLA 优化）下的时间线片段。


图 1：TensorFlow 计算 GELU 的时间线
其中，黄色矩形框中对应的是激活函数 GELU。可以看到，在 TensorFlow 中，这个函数是通过 8 个类似 Pow，Add，和 Tanh 等基本 OP 来实现的。Layer Normalization 操作也是类似的情况：


图：TensorFlow 计算 Layer Normalization 的时间线
在 TensorFlow 中，每一个基本 OP 都会对应一次 GPU kernel 的调用，和多次显存读写，这些都会增加大量额外的开销。TensorFlow XLA 可以在一定程度上缓解这个问题，它会对一些基本的 OP 进行合并，以减少 GPU kernel 的调度和显存读写。但在大多数情况下，XLA 依然无法达到最优的性能，特别是对于 BERT 这种计算密集的情况，任何性能的提升都将节省巨量的计算资源。 如英伟达计算团队前面提到的，OP 融合可以降低 GPU 调度和显存读写，进而提升性能。出于性能最大化的考虑，在 Faster Transformer 内部，开发团队将除矩阵乘法以外的所有 kernel 都进行了尽可能的融合，单层 Transformer 的计算流程如下图所示：


图：BERT 中 Transformer Layer 的计算流程图
如上图所示，Faster Transformer 只用了 14 个 kernel 就完成了原来将近 60 个 kernel 的计算逻辑。这其中，8 个 kernel 是通过调用 cuBLAS 接口计算矩阵乘法（绿色框），其余 6 个是自定义 kernel（蓝色框）。 针对 batch size 比较小的场景（例如问答，TTS 等），简单的融合后，基本上就可以达到很好的性能。这类场景下，TensorFlow 原生实现的最大瓶颈就在于频繁的 kernel launch，融合后大大降低了 launch 的开销，因此可以比较轻易地获得很好的加速效果。 针对大 batch 的场景，我们需要对矩阵乘法和所有的自定义 kernel 做精细的调优，才能达到很好的加速效果。英伟达计算团队从矩阵乘法算法选择，非矩阵乘法操作的参数配置，SoftMax 多版本实现，以及数据结构类型等几个方面对大 batch 的情况进行了专门的调优。 首先针对矩阵乘法，在调用 cuBLAS 的接口时，可以指定性能最优的算法。特别是针对 Volta 和 Turing 架构的 GPU，使用 Tensor Core 进行半精度计算时，当精度满足需求的情况下，累加器也可以选择半精度，从而进一步提升性能。 除矩阵乘法以外的 6 个 kernel，大部分都是对矩阵乘的结果进行一些 element-wise 的操作。输入矩阵的大小，跟 4 个参数有关，batch size，句子长度，attention 的 head 数量以及每个 head 的维度。针对不同的应用场景，参数大小可能极为不同。比如在线问答类的场景，batch size 可能为会很小，通常为 1。
而广告推荐或者搜索类的场景，batch size 通常跟候选集大小有关，一般会是几百的规模。这样，输入矩阵的行数变化范围可能是几十到上千。因此，我们需要针对不同的情况，动态的调整 kernel launch 时的配置参数（grid 和 block 的大小），甚至要针对同一个功能实现多个不同版本的 kernel 函数，例如，SoftMax 的计算就有两个不同的实现版本。 针对半精度 FP16，英伟达计算团队对各个 kernel 也进行了相应优化。首先，在 kernel 的实现中，将输入的 half 指针转成 half2 类型，并使用了 half2 相关的数学函数。这样不仅仅可以达到 2 倍于 half 的访存带宽和计算吞吐，还可以极大地减少指令的发射数量。其次，在 SoftMax 以及 Layer Normalization 的操作中，为防止求和溢出，将数据以 half2 的形式读入后，会转成 float2 类型，来做求和计算。 除上述优化之外，Faster Transformer 还优化了前向计算中耗时较高的 GELU 激活函数，Layer Normalization 以及 SoftMax 等操作。比如利用 warp shuffle 实现高效的矩阵按行求和操作，将 1/sqrtf 计算替换为 rsqrtf 函数，以及 power (x, 3.0) 替换为 x * x * x 等。总之，英伟达计算团队针对 Transformer 进行了各种优化以保证它的高效执行。



WAIC 2019 开发者日将于 8 月 31 日在上海世博中心举办，包含 1 个主单元、4 个分单元、黑客马拉松比赛和开发者诊所互动区。

届时，全球顶尖 AI 专家、技术大牛、知名企业代表以及数千名开发者将齐聚上海，围绕人工智能前沿理论技术和开发实践进行分享与解读。

点击阅读原文，立即报名。