英伟达公开课 | 如何突破Decoder性能瓶颈?揭秘FasterTransformer2.0的原理与应用,已开源

位来 发自 凹非寺

量子位 报道 | 公众号 QbitAI

4月9日,英伟达x量子位分享了一期nlp线上课程,来自NVIDIA的GPU计算专家、FasterTransformer 2.0开发者之一的薛博阳老师,与数百位开发者共同探讨了:

  • FasterTransformer 2.0 新增功能介绍

  • 如何针对Decoder和Decoding进行优化

  • 如何使用Decoder和Decoding

  • Decoder和Decoding能够带来什么样的加速效果

应读者要求,我们将分享内容整理出来,与大家一起学习。文末附有本次直播回放、PPT链接,大家也可直接观看。

以下为本次分享的内容整理:

大家好,今天为大家介绍的是 FasterTransformer 2.0的原理与应用

什么是FasterTransformer?

首先,参加本期直播的朋友对Transformer架构应该都有了一定了解。这个架构是在“Attention is All You Need”这篇论文中提出的。在BERT Encoder中使用了大量Transformer,效果很好。因此,Transformer已成为 NLP 领域中非常热门的深度学习网络架构。

但是,Transformer计算量通常是非常大的。因此,Transformer的时延往往难以满足实际应用的需求。

Attention is All You Need截图

Transformer架构可以应用于Encoder或Decoder上。在Encoder中,Transformer包含1个multi-head attention和1个feed forward网络,在Decoder中,包含2个multi-head attention和1个feed forward网络。

其中,纯Encoder架构在目前的很多应用中都有很好的表现,比如Q&A系统、广告推荐系统等,因此,针对Encoder的优化是非常有必要的。

而在一些场景中,如翻译场景,我们需要Encoder和Decoder架构。在这种架构下,Decoder消耗的时间占比是非常高的,可能达到90%以上,是推理的主要瓶颈。因此,针对Decoder的优化也是一项重要的工作,能带来明显的加速效果。

实际应用中,FasterTransformer 1.0版本针对BERT中的Encoder为目标做了很多优化和加速。在2.0版本中,则主要新增了针对Decoder的优化,其优越的性能将助力于翻译、对话机器人、文字补全修正等多种生成式的场景。

上表比较了Encoder和Decoder计算量的不同。当我们需要编解码一个句子的时候,Encoder可以同时编码很多个字,甚至可以直接编码一个句子。

但是Decoder是一个解码的过程,每次只能解码一个字,因此,解码一个句子时我们需要多次Decoder的forward,对GPU更不友善。

Faster Transformer框架

上图列出了FasterTransformer中针对BERT优化的模块。在编码方面,以BERT为基准,提供了一个单层的、等价于BERT Transformer 的模块,供使用者做调用。当我们需要多层的Transformer时,只需调用多次Encoder即可。

解码方面更为复杂,为了兼顾灵活性与效率,我们提供两个不同大小和效果的模块:

Decoder(黄色区块) 由单层的 Transformer layer 组成,它包含两个attention和一个feed forward 网络;而Decoding(蓝色区块)除了包含多层的 Transformer layer 之外,还包括了其他函数,例如 embedding_lookup、beam search、position Encoding 等等。

我们用一个简单的虚拟码展示Decoder和Decoding的区别。

在Decoding中通常有两个终止条件,一是是否达到预先设定的最大的sequence length,第二个条件是所有的句子是否都已经翻译完毕,未终止时会不断循环。

以句子长度为128的句子翻译场景为例,若其 Decoder 是由6层的 Transformer layer 组成的,总共需要调用 128×6=768 次的Decoder;如果是使用 Decoding 的话,则只需要调用一次Decoding,因此Decoding的推理效率更高。

小结

首先,FasterTransformer提供了高度优化过的Transformer layer:在Encoder方面是基于BERT实现的;在Decoder方面基于OpenNMT-TensorFlow开源的库做为标准;Decoding包含了翻译的整个流程,也是基于OpenNMT-TensorFlow。

其次,FasterTransformer 2.0的底层由CUDA和cuBLAS实现,支持FP16 和 FP32 两种计算模式,目前提供C++ API和TF OP。

现在,FasterTransformer 2.0已经开源,大家可以在 https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer/v2 获取全部源代码。

如何进行优化?

先以Encoder为例。

TF Encoder Transformer layer

参数:no XLA,batch size 1,12 heads,size per head 64,FP 32

图中蓝色方块表示GPU在实际运行,空白的表示GPU在闲置,因此GPU在很多时间是闲置状态。造成GPU闲置的原因是kernels太小,GPU要不断闲置以等待CPU启动kernel的时间,这也称为kernel launch bound问题。

如何解决这个问题?

我们尝试开启TF的XLA,其他参数不变。图中我们看到,从原本计算1层Transformer layer需要50个kernel缩减到24个左右。大部分kernel变得比较宽,虽有加速,但是闲置的时间还是比较多。

因此,我们提出FasterTransformer针对Encoder进行优化。

首先,我们把矩阵计算部分挑选出来,用NVIDIA高度优化的库cuBLAS 来计算,此外的部分,我们把能融合的kernel都尽可能融合起来。

最终的结果如上图右边,经过整体的优化后,我们只需要8个矩阵计算加6个kernel就可以完成单层Transformer layer计算,也就是说所需kernel从24个减少到14个。

我们可以看到,优化后每一个 kernel 都相对比较大,时间占比小的kernel也减少了。但还是有很多空白的片段。

我们直接调用C++ API,如图,GPU闲置的时间几乎没有了。因此,小batch size情况下,我们推荐使用C++ API获得更快的速度。 当batch size比较大时,GPU闲置时间会比较少。

接下来我们看下Decoder。

参数:no XLA,batch size 1,8 heads,size per head 64,FP32

经过统计,TF需要使用70个左右kernel来计算1层Transformer layer。直观来看,非常小、时间占比非常短的kernel更多。因此,batch size比较小的情况下,优化效果会更明显。

Decoder的优化同上述Encoder,特别之处是,Decoder里面的矩阵计算量非常少,因此我们把整个multi-head attention以一个kernel来完成。经过优化之后,原本需要70个kernel才能完成的计算,只需要使用16个kernel就能够完成。

在更大的Decoding模块中,另一个时间占比较多的kernel是beam search,这里我们针对top k做出优化。在GPU中可以同时执行多个block和多个warp,并行运算,大大节省时间。

更多优化细节

如何使用FasterTransformer?

大家可以在 https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer/v2 根目录下找到对应资料:

针对 Decoder 和 Decoding,FasterTransformer 分别提供了 C++ 和 TensorFlow OP 这两种接口。

C++接口

首先创建一个Eecoder,超参数如图:

其次,设定训练好的模型权重;

设置好后,直接调用forward即可。

TF OP接口

首先,我们需要先载入OP。这里以Decoder为例,会自动创建TF需要使用的库,调用接口时先导入.so文件(图中已标红):

然后调用Decoder,放入input、权重、超参数,然后针对out put 做Session run。

这里需要注意的是,参数里有一个虚拟的输入 (pseudo input)。这个输入是为了避免 TensorFlow 的 decoder 和 FasterTransformer Decoder 发生并行,因为我们发现并行执行时,Decoder中的memory可能会被污染。实际应用的时候可以将这个输入拿掉。

优化效果

最后我们来看下优化的效果如何。首先测试环境设置:

使用的GPU是NVIDIA的Tesla T4和V100。

Encoder模块在Tesla V100的结果

超参数设置:12 layers,32 sequence length,12 heads,64 size per head(BERT base),under FP 16

结果如上图,batch size从100逐步增加到500的过程中,FasterTransformer对比TF开启XLA,大概可以提供1.4倍的加速。

Decoder和Decoding模块在Tesla T4的结果

超参数设置:Batch size 1,beam width 4,8 heads,64 size per head,6 layers,vocabulary size 30000,FP 32

结果如上图,不同的sequence length下,相比于TF,FasterTransformer Decoder可以带来3.4倍左右的加速效果,Decoding可以带来7-8倍的加速,效率更高。

超参数设置:Batch size 256,sequence length 32,beam width 4,8 heads,64 size per head,6 layers,vocabulary size 30000

结果如上图,把batch size固定在较高值时,不同的FP下,FasterTransformer Decoder和Decoding也带来一定的加速效果。

最后,本次直播的PPT获取连接: https://images.nvidia.cn/cn/webinars/2020/apr09/20200409_FasterTransformer.pdf

本次直播回放:

如何关注、学习、用好人工智能? 

每个工作日, 量子位 AI内参 精选全球科技和研究最新动态,汇总新技术、新产品和新应用,梳理当日最热行业趋势和政策,搜索有价值的论文、教程、研究等。

同时,AI内参群为大家提供了交流和分享的平台,更好地满足大家获取AI资讯、学习AI技术的需求。扫码即可订阅:

了解AI发展现状,抓住行业发展机遇

AI社群 | 与优秀的人交流

量子位  QbitAI · 头条号签约作者

վ’ᴗ’ ի 追踪AI技术和产品新动态

喜欢就点「在看」吧 !