你见过 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 进行了各种优化以保证它的高效执行。