『』微信也在用的Transformer加速推理工具，现在腾讯开源了( 二 )

本文插图
Turbo技术原理
能够取得如此好的推理性能，这背后的计算原理又是什么呢？
TurboTransformers的软件架构如下图，它让微信内部众多NLP线上应用能够充分榨取底层硬件的计算能力，让算法更好地服务的用户。
具体来说TurboTransformers可以在算子优化、框架优化和接口部署方式简化三个方面做了工作。

本文插图
算子层优化
Transformer都包含了什么计算呢？
如下图所示，图(a)展示了论文Transformer结构示意图，这里称灰色方框内的结构为一个Transformer Cell ， BERT encoder堆叠了Nx个这样的Transformer Cell 。
图(b)将一个Cell的细节加以展开，每一个矩形都是一个独立的计算核心。

本文插图
Transformer Cell计算包含了8个GEMM(通用矩阵乘法， General Matrix Multiplication)运算。通过调优Intel MKL和cuBLAS的GEMM调用方式来获得最佳GEMM性能。
并且在硬件允许条件下，在GPU上使用tensor core方式进行GEMM运算。
类似NVIDIA FasterTransformers方案，将所有GEMM运算之间的计算融合成一个调用核心。融合会带来两个好处，一是减少了内存访问开销，二是减少多线程启动开销。
对于这些核心，在CPU上采用openmp进行并行，在GPU上使用CUDA进行优化实现。
对于比较复杂的LayerNorm和Softmax算子，它们包含了不适合GPU上并行的规约操作， TurboTransformers为它们设计了创新并行算法，极大降低了这些算子的延迟。
理论上Transformers推理延迟应该近似于矩阵乘法延迟。
框架层优化
TurboTransformers采用了一个有效的内存管理方式。
由于NLP的采用变长输入特性，每次运算中间结果的大小其实并不相同。为了避免每次都分配释放内存，研究人员通过Caching方式管理显存。
为了能够无缝支持pytorch/tensorflow训练好的序列化模型，提供了一些脚本可以将二者的预训练模型转化为npz格式，供TurboTransformers读入。
特别的，考虑到pytorch huggingface/transformers是目前最流行的transformers训练方法，支持直接读入huggingface/transformers预训练模型。
应用部署
Turbo提供了C++和Python调用接口，可以嵌入到C++多线程后台服务流程中，也可以加入到pytorch服务流程中。
研究人员建议TurboTransformers通过docker部署，一方面保证了编译的可移植性，另一方面也可以无缝应用于K8S等线上部署平台。
传送门
GitHub项目地址：https://github.com/Tencent/TurboTransformers/blob/master/README_cn.md
— 完 —
量子位 QbitAI · ***签约
关注我们，第一时间获知前沿科技动态