高性能Transformer推理分析
发布时间:2022-10-17 07:00
摘 要
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
本文中所引用的图片和视频来源均已标明,如有侵权请联系删除。
上一个: 由简入繁探究机器视觉中的数据增强(下)
近期文章
通用AI模型的未来:深度强化学习(deep reinforcement learning)
2023-05-08
2023-04-24