摘  要

 

Transformer已经成为了目前非常具有影响力的AI模型，而如何将其快速地部署到云端或者本地端进行高效地推理，就成了AI落地过程中非常重要的问题。相比于训练过程，在保证适配性的同时，推理过程更加重视模型的计算延迟和吞吐性能。为了应对这些问题，近期出现了多种以C++/CUDA方式来实现的Transformer推理加速工具。本文关注了这些工具中用到的一些加速方案，如算子融合、低精度加速、矩阵乘法的张量加速、多卡并行等。为了能够详述其中一些细节，这里以FasterTransformer库为例，介绍其在GPT模型推理上的实现细节。希望本文能够帮助读者进一步了解和使用这些Transformer推理加速工具。

 

引  言

 

在自然语言处理任务上的成功，正使得Transformer成为具有超高影响力的AI模型。相比于卷积网络和RNN网络，Transformer对长距离上下文特征的捕捉能力更强。它已经从最初的自然语言翻译、问答和情感分析等任务的应用，逐渐扩展到大部分常见的AI 任务，如计算机视觉、语音识别，甚至是分子结构分类等。Transformer模型本身采用的是编码器—解码器结构。编码器将时序序列映射到隐藏层，而解码器将隐藏层再映射为时序序列，从而使其可以应用于下游任务。

 

近年来，这些应用于不同领域的Transformer模型，正以预训练-微调（pretrain-finetune）或者以预训练-提示（pretrain-prompt）方式快速地发展。通常的做法是在大型语言或者图像数据集上对模型进行预训练，然后根据特定的任务如情感分析、图像分割等下游任务，对已训练的网络权重进行部分微调后推理，或者直接使用提示符进行推理。而这种预训练再加微调或者提示的大模型使用方式，也极大地提高了模型的可复用能力。且由于越大的模型将会有更好的效果表现，因此也推动了大模型的快速流行，例如Open AI提出的GPT-3模型已经达到了1750亿参数 [1]。

 

然而如何应用这些已训练好的大模型，如何将其快速地部署到云端或者本地端从而方便用户使用，就成了AI落地过程中非常重要的问题。相比于以往的卷积网络，Transformer模型更加依赖大算力的支撑。例如，对于40个字的文本序列，进行一次Bert推理需要 7 Gflops，由中文翻译到英文的Seq2Seq模型需要 20 Gflops，而卷积网络如ResNet50对 3*224*224 的图像只需要更小的 4 Gflops 算力[2]。

 

由于Transformer模型通常由深度学习框架如PyTorch或者TensorFlow训练得到，那么一个自然的问题是，为何不直接使用原生的深度学习框架进行推理呢？这其中的主要原因是为了获得更好的性能和灵活性。AI模型在训练和推理过程有不同的需求，原生的深度学习框架，为了更好地支持如反向传播等这些训练过程，提供了更多的通用性同时，也牺牲了推理过程中的性能。原生模型往往进行了更细粒度算子拆分，从而导致额外的计算，以及过量的显存分配，因此其无法在推理过程中真正发挥底层硬件的性能。

 

为了更大程度地提高Transformer在推理过程中的性能，发挥GPU硬件的性能，近期出现了一些用C++/CUDA方式来实现Transformer推理的工具。本文的主要目的是梳理近期的一些进展，介绍其在推理上的性能优化思路，如算子融合，低精度加速，Tensor core 加速，多卡并行等。为了能够详述一些细节，本文以FasterTransformer库，以及GPT模型的推理过程为例进行介绍，从而帮助读者了解和使用这些推理加速工具。

 

Transformer推理库

 

为了能够规模化地将Transformer部署在带有GPU的高性能数据中心上，总体上需要针对两个方面进行优化：首先是高效性，即在线的推理服务能够具有更低的延迟，以及更高的吞吐；另一方面是灵活性，针对多种自然语言处理任务能够灵活地适应于变长的输入序列情形。尤其是第二点，虽然目前市面上也有如Tensorflow XLA [3]、TVM [4]、tensorRT [5] 这样的工具，但是这些工具的主要问题是无法很好地适应于可变输入长度的问题，如自然语言处理问题中长短不一的句子或语音序列。尽管像onnxruntime [6] 也有一些适应于可变输入长度的能力，但是对于GPU内存的使用并不高效，也无法及时地释放显存。

 

目前针对以上两个方面的需求，出现了一些为此适配的Transformer推理库，如英伟达的FasterTransformer [7]，腾讯的TurboTransformers [8]，字节跳动的LightSeq（LightSeq2在LightSeq基础上增加了训练功能）[9]，网易伏羲的EET [10]，另外微软的DeepSpeed（致力于提高模型训练的性能的库）[11] 也包含了推理过程的优化，以及针对OpenNMT翻译模型的CTranslate2 [12]和专门用于提高T5模型推理速度的Fast T5 [13]。同时这些库还额外集成了一些易用的功能，如与PyTorch、Tensorflow进行了插件式集成，以及封装了在线服务等。

 

优化方案

 

以上这些推理库主要由C++/CUDA编写，并使用Python进行功能的封装，总体上说其性能优化的思路可以有：算子融合、低精度优化、Tensor Core加速、张量并行与流水并行、缓存计算结果等。下面来看一下这些Transformer推理库用到的一些优化方案。

 

1. 算子融合

 

在深度学习框架中，复杂的计算被拆分成了细粒度的算子，通过拼合细粒度算子就可以得到复杂的算子，这样就使得深度学习框架具有了更好的通用性。相反地，为了获得更好的性能，则需要将细粒度算子拼合成一个CUDA核函数，从而减少GPU核函数的调用以及显存分配过程。如下图，例如其中的QKV三个矩阵乘法（GEMM）计算过程都被融合到了一个gemm0123中，而add bias以及Layer Norm的计算过程则被合并到了一个算子中，从而达到了性能优化的过程。

 

图1.Transformer解码层核函数的算子的融合对比（右图的融合算子图）[2]

图片来源：https://arxiv.org/pdf/2010.05680.pdf

​

2. 低精度加速

 

在网络推理过程中，默认的推理精度为FP32，而当使用FP16时，甚至是使用INT8进行量化推理时，将得到很大的性能优势。下图是FasterTransformer对比FP16和IN8得到的性能提升比例，可以看到对于大批次长句子时，使用INT8-v2的量化方式，可以得到5倍的性能优势 [7]。

 

图2. 不同批大小和序列长度下的FP16与INT8加速对比 [7]

图片来源https://github.com/NVIDIA/FasterTransformer/blob/main/docs/images/Py_Encoder_T4.png
 

3.Tensor Core加速

 

在Transformer模型的计算过程中，矩阵乘法（GEMM）是整体的计算瓶颈，如下图LightSeq论文的实验结果，因此GEMM的优化将会起主要作用。而GPU的Tensor Core是专门针对GEMM的加速单元，因此在这些推理库中都会调用cuBLAS来启用GPU的Tensor core来进行加速。注意在Transformer模型的推理中，编码器相比于解码器的计算更依赖GEMM的计算。

 

图 3 LightSeq库中FP16与FP32下计算占比 [14]

图片来源：https://arxiv.org/pdf/2010.13887.pdf

 

4. 张量并行与流水并行

 

FasterTransformer在4.0版本时引入了并行解码过程，可以将Transformer模型通过张量并行以及流水并行的方式，即将计算拆分到多张的GPU卡上进行，使得多张GPU同时计算来降低模型计算的延迟。

 

图 4 张量并行与流水并行[15]

图片来源：https://developer.nvidia.com/zh-cn/blog/accelerated-inference-for-large-transformer-models-using-nvidia-fastertransformer-and-nvidia-triton-inference-server/

5. 缓存计算结果

 

在GPT推理过程中，即对文本进行逐个单词生成时，无需每次都将整个QKV矩阵输入到多头注意力层来进行计算，而是可以将先前计算的K和V值结果进行缓存，只需要增量式地放入新的单词的Q值来进行计算。

 

图 5 缓存GPT上一次推理的计算结果（右图为缓存计算结果）[16]

图片来源：https://huggingface.co/blog/accelerated-inference
 

 

6. 其他方法

 

除了以上提到的一些优化方案，在采样时还可以使用TopK（候选字符集有K个），TopP（候选字符集的总概率达到P）的方式，还有如诸如束搜索方法（beam search），以及可通过早停的方式使得小概率的分支尽快停止。另外可以针对一些重点算子如softmax、gelu、layernorm进行手工优化，或是针对不同的批大小（batch size）和长序列问题进行单独地优化等。

 

FasterTransformer库

 

为了更好叙述这些库的细节，这里以FasterTransformer (FT)库为例。FT库从2019年7月发布1.0版本以来，到2022年八月已经发布到5.1版本，中间迭代多回，在性能与功能上都有了一定的提升。其他推理库也都有相应的各自特色，例如在LightSeq [14]以及EET [17] 的论文中，都对比了某个版本的FT库，得到了更进一步的性能提升。

 

FT库是用C++/CUDA编写的，其中矩阵乘法部分调用了高性能矩阵计算库cuBLAS、cuBLASLt、cuSPARSELt 库。而FT库针对编码器-解码器结构分别进行了优化，用户可以单独使用解码器部分（如BERT），或者仅解码器部分（如 GPT），或者两部分都有的结构（如T5）进行推理。根据英伟达的显卡架构版本，加速方案包括FP16，INT8低精度加速，稀疏矩阵乘法优化，以及分布式计算，即通过张量并行和流水并行将大规模神经网络拆分到多个GPU进行推理。

 

整个项目为Python做了接口的封装，主要功能由C++实现，C++项目的代码由CMake来进行组织。为了更加具体地展示FT库的结构特色和其功能的具体实现，这里以GPT模型为例。从文件目录来分，主要分为整体模型（models/gpt）、模型层(layers)、核函数(kernels)、辅助函数（utils）、扩展接口（Tensorflow/PyTorch/Triton的算子扩展）。

 

GPT的应用启动入口在examples/cpp/gpt/gpt_example.cc，主程序总体上依次实现了三个过程：1）载入参数配置信息；2）载入已训练模型的网络权重；3）持续地通过GptDecoder来生成字符。其配置文件在同一目录下的gpt_config.ini，里面包含有批处理大小、输出尺寸、注意力头的个数、模型权重文件位置等可调整的配置信息。其具体的计算过程如下图：

图 6 GPT的计算流程[18]

图片来源：https://github.com/NVIDIA/FasterTransformer/blob/main/docs/images/gpt/gpt.png

 

计算流程是:输入一个批次的序列，之后经过Embedding层之后，输入到GptDecoderLayer进行多层解码，之后再通过层会归一化（LayerNorm）和矩阵乘法（GEMM），输出时会进行topK/topP采样或者通过束搜索（beam search）的方式进行。其中每一层的GptDecoderLayer又包含了层归一化、自注意力等融合的算子，全连接层又可分解为矩阵乘法和激活层。为了看清楚各部分的计算，图中对其进一步进行了标示。蓝色部分为核函数的计算，绿色处为模型层的计算，其实现代码在layers目录下，而蓝色处为矩阵乘法，相应的矩阵乘法可以通过GPU上的Tensor core单元进行加速。由于实现上细节比较多，这里不再一一展开叙述。读者可以在代码库中看到上文提到的各种优化方案的具体实现。

 

小  结

 

本文主要关注了近期Transformer推理加速的相关进展。为了更大程度地发挥GPU硬件的性能，以及保持对可变长序列的通用性，近期出现了多种以C++/CUDA方式实现的Transformer推理的加速工具。这些工具在实现加速时都有一些共性的优化方案，例如小算子融合为大算子核函数，用低精度FP16、INT8来进行加速，使用GPU上的Tensor core单元对GEMM进行加速，以及使用多张GPU卡来并行计算等。为了能够详述一些细节，这里进一步介绍了FasterTransformer库对于GPT模型的实现细节，从而使读者能够进一步了解这一类工具。文中存在不足的地方请各位读者批评指正，也欢迎大家一起参与我们的讨论。

 

参考文献

 

1. Floridi, Luciano, and Massimo Chiriatti. "GPT-3: Its nature, scope, limits, and consequences." Minds and Machines 30.4 (2020): 681-694.

2. Fang, Jiarui, et al. " TurboTransformers: an efficient GPU serving system for transformer models." Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. 2021.

3. https://www.tensorflow.org/xla.

4. https://github.com/apache/tvm.

5. https://github.com/NVIDIA/TensorRT.

6. https://github.com/microsoft/onnxruntime.

7. https://github.com/NVIDIA/FasterTransformer.

8. https://github.com/Tencent/TurboTransformers.

9. https://github.com/bytedance/lightseq.

10. https://github.com/NetEase-FuXi/EET.

11. https://github.com/microsoft/DeepSpeed.

12. https://github.com/OpenNMT/CTranslate2.

13. https://github.com/Ki6an/fastT5

14. Wang, Xiaohui, et al. "Lightseq: A high performance inference library for transformers." arXiv preprint arXiv:2010.13887 (2020).

15. https://developer.nvidia.com/zh-cn/blog/accelerated-inference-for-large-transformer-models-using-nvidia-fastertransformer-and-nvidia-triton-inference-server

16. https://huggingface.co/blog/accelerated-inference

17. Li, Gongzheng, et al. "Easy and Efficient Transformer: Scalable Inference Solution For large NLP model." arXiv preprint arXiv:2104.12470 (2021).

18. https://github.com/NVIDIA/FasterTransformer/blob/main/docs/gpt_guide.md

本文中所引用的图片和视频来源均已标明，如有侵权请联系删除。