GPU BERT上线性能不合格，看看微信AI的PPoPP论文

以BERT为代表的Transformer模型的出现是自然语言处理（NLP）领域近年来最关键的算法创新。目前互联网公司内很多线上服务系统，比如推荐系统、对话系统、翻译系统等，都使用了Transformer模型。但是由于Transformer巨大的计算量和变尺度输入的特性，用现有的解决方案部署线上服务，效果往往不能达到预期。

微信AI近日的一篇论文解决了这个问题。本文基于论文《TurboTransformers: An Efficient GPU Serving System For Transformer Models》，该论文即将发表在ACM 并行编程原理与实践会议 (ACMSIGPLAN Symposium on Principles and Practice of parallel Programming, PPoPP 21') 上。

录取的论文介绍了对目前NLP领域流行的Transformer结构模型的一系列优化方法，包括GPU计算效率、内存管理和服务调度优化，最终使得Transformer线上服务获得超过PyTorch，TensorFlow，onnxruntime，TensorRT等现有方案的业界最佳性能。文章中介绍的技巧已经应用于模式识别中心2020年4月份对外开源项目TurboTransformer之中，它也是腾讯开源协同的成果之一。下面我们对论文一些细节一睹为快吧。

问题背景

与递归神经网络（RNN）模型不同，Transformer对长序列的不同位置信息进行并行处理，因此可以显著提高处理长序列任务的准确性。但是，要在GPU数据中心的在线服务上有效部署它们并非易事。暂且不说Transformer结构相比RNN引入了的更多计算量，Transformer首次将变尺寸输入问题引入到了线上推理服务工程优化之中。

在NLP领域，推理任务的输入常常是长度不定的句子。与之相反，在计算机视觉（CV）领域，推理任务的输入一般是固定尺寸的图片。尽管RNN的模型（如LSTM和GRU）也具有处理变长输入的能力，但它们是通过将变长输入切分若干固定长度输入后按顺序串行执行，这相当于将变长输入任务转换为定长输入任务。Transformer模型则对输入序列维度进行并行处理，这正是它带来更好模型表现的原因。

尔之蜜糖，我之砒霜，由于无法预知输入序列长度，Transformer算法特点给其线上服务的工程实现带来巨大挑战。首先，传统适合CV推理的解决方案，比如TF-Serving，面对新的问题则捉襟见肘：它们往往需要预置输入尺寸，然后运行离线的前处理对计算方式进行调优，使用它们部署Transformer服务时，需要对输入进行补零成预置的尺寸，这造成大量计算资源浪费。其次，服务框架的优化变得更加复杂。TF-serving利用批处理技术（Batching）来提高GPU执行效率。在一批可变长度的请求中，短请求的补零会带来额外的计算量，这通常与批处理带来的性能提升的作用相冲突。

为Transformer搭建GPU服务仍然缺少一套最佳实践方式，为解决此问题，微信AI研发了一套名为TurboTransformer的解决方案。三个创新设计使它在同类产品中脱颖而出。

第一，设计了更高效的GPU的批量规约算法，可以应用于Transformer中非矩阵乘法的主要热点Softmax和LayerNorm核心。

第二，针对可变长度输入情况设计了一种内存分配算法，该算法可以更好地平衡内存占用量和分配/释放效率。

第三，使用动态规划算法设计了新的批处理策略的服务框架，针对可变长度服务处理请求，可以达到理论最佳吞吐量。该系统可以在GPU平台上达到目前最佳的Transformer模型服务性能，并且可以通过几行代码无缝地集成到PyTorch代码中。

论文中TurboTransformer推理系统由推理runtime和服务框架两部分组成，目前图1的runtime部分已经在github上开源：https://github.com/Tencent/TurboTransformers

图1，TurboTransformer线上推理系统

高效的批量规约化方法

TurboTransformer采用了和NVIDIA FasterTransformer类似的核心融合方案来避免非矩阵乘法CUDA核心启动开销，增加显存数据的局部性。在此基础上，它进一步优化了Softmax和LayerNorm的计算效率。这两个算子都可以看成一批一维数组规约（如求和，求最值）操作，因此被统一称为称为批量规约操作。

下图的顶部显示了Faster Transformers中对批量规约操作的经典实现。该工作借鉴了GPU上一维数组规约操作的最佳实践，将n个一维规约负载（在图顶部的虚线内）分配给CUDA线程块，线程块顺序执行n次规约操作。一次规约操作通过两次遍历完成，每次遍历利用warpReduce操作对32个数据进行规约，两次遍历之间需要一次线程块内的全局线程同步。

图2，本文提出的并行批规约化方法和传统方案的比较。

在上述算法中，仍留有一些空间用于提升warp级别和指令并行性。首先，由于共享内存的访问，线程同步带来了巨大的开销。其次，如果输入数组未对齐32，则边界处理导致的warp执行不同逻辑，从而引入额外的分支判断开销。第三，warpReduce内存的指令发射效率低下。在图2的右上角，SHFL DOWN指令中的目标寄存器R3作为FADD指令中的源寄存器。只有在SHFL完全完成之前才发出FADD指令，其等待时间超过1个周期会造成指令流水线的卡顿。

TurboTransformer在GPU上突破了传统的批量规约算法的效率边界，它通过多次一维归约操作之间的并行化来克服以上三个缺点。如图2的底部所示，引入了一个新的子例程warpAllReduceSum_XElem（在图中X = 2），该子例程将X个warp进行批处理，一起进行规约。首先，blockReduceSum_XElem中规约X行元素只需要一次同步，因此减少了（X − 1）/ X同步成本。其次，可以将X独立边界处理合并为一个，从而减少warp分支。第三，因为消除了指令依赖性，指令发布效率更高。如图所示，FADD需要在两个周期之后将SHFL DOWN的目标寄存器用作源寄存器，此时，不依赖于前一个的另一个SHFL DOWN可以立即发射。

可感知序列长度的内存分配算法

除了计算层面的优化，内存的优化对推理服务性能提升也至关重要。由于输入尺寸在服务中并不是确定的，神经网络每层的输出张量的内存使用时刻变化。GPU上分配释放内存开销很大，如不加优化会造成计算等待内存分配的现象。目前常见的GPU内存分配算法将已经分配内存缓存起来供下次分配使用，但是这样往往造成推理runtime占据大量内存，影响数据中心GPU使用效率，并限制了模型加载数量和尺寸。

为了获得更高的分配效率和更少的占用空间，TurboTransformer对变化的中间张量采用了一种创新的内存优化方法。该方法在获得每次推理的输入后就启动一个轻量级的可变长度感知内存管理器，内存管理算法结合了缓存和计算图拓扑结构信息，在O(N^2)复杂度时间内获得最优的内存分配方案。

首先，它以块为粒度（如2MB）分配释放内存。不仅避免频繁分配小块内存，通过重用已经分配的块，分配效率可以在服务期间保持较高水平。其次，它利用计算图拓扑结构预先知道每个中间张量的生命周期，并在识别到新到达请求的序列长度后立即计算特定块内每个张量的在内存块中的偏移量。这样，没有重叠生命周期的张量可以重用相同的内存空间，因此可以尽可能减少内存占用。

下图显示了将这个的算法应用于BERT推理应用程序的示例。当输入长度从200变为240时，只需要再分配一个块，并调整每个中间张量在内存块中的偏移量。

图3，使用了本文提出的可变长度感知分配算法的一个内存分配示例。

可感知序列长度的服务批处理算法

使用runtime可以非常方便地搭建一个线上服务，但使之满足一定延迟和吞吐量标准却并不容易。用户请求一般都是batch=1的输入序列，小批量（batch）会导致GPU硬件利用率低。将多个请求打包成一个相对较大的批处理并一起进行推断可以提高硬件利用率，这也是TF-Serving为代表服务框架的重要优化方案。但是，正如我们提到的，变尺度输入使得传统批处理优化方法并不适用。

如下图所示，如果按照传统批处理方法，这五个请求都应该被打包成长为77的序列，这样有大量补零开销。实际上，获得最优吞吐量的批处理方法如图右所示。

图4，可变长度请求的批处理调度方法的示例。

论文中提出了一种创新的序列长度感知批处理调度算法。核心思想是使用动态规划（DP）在 O（n^2）时间复杂度内得到最大的响应吞吐量的目标。该算法的输入由两个数据结构组成。request_list是具有可变长度的输入请求的列表。cached_cost是一个字典。它使用两个关键字作为索引，即序列长度和batch大小。其值是使用相应键为参数的推理成本。服务首次在特定硬件上启动后，将通过预热阶段收集cached_cost的值，该硬件利用运行时在所有可能的批处理大小和序列长度下运行推断。它们存储在磁盘或数据库（图1中的DateBase）中，并在重新启动服务模块时重新加载到内存中。首先，request_list根据序列长度按升序排序。然后，初始化一个状态的数组state，以存储中间信息。具体地，state[i]记录用于处理子列表request_list [0：i]的最小时间开销。该算法遍历request_list中的每个请求，并使用DP算法更新相应位置i处的state[i]。如下显示了该DP问题的Bellman方程。

实验效果

TurboTransformers的运行时显示出明显的性能优势。在变长随机输入的测试下，将运TurboTransformer行时与PyTorch（版本1.5.0）进行比较。对于Bert，同时也与onnxruntime-gpu（版本1.3.0）运行的结果相比较。对于Bert，Turbo到PyTorch的加速范围为1.10x-2.58x；Turbo到onnxruntime的加速范围为0.84x-1.68x。对Albert，Turbo到PyTorch的加速范围为1.35x-2.26x。在一个翻译模型中，Turbo到PyTorch的加速范围为1.85x-2.51x。

图5，可变长度请求上的三个模型的运行时的延迟。

论文同时也测试了定长输入情况下TurboTransformer runtime性能，比较对象包括PyTorch 1.5.0、Faster Transformer V1，Onnxruntime 1.3.0，TensorFlow 1.13，TensorRT 5.1.5等一些不具备处理变长输入序列能力的方案。下图显示了RTX 2060和V100上的TurboTransformers对其他方案的加速。在RTX 2060 GPU上，除了最低工作负载情况的个别情况之外，TurboTransformer的运行时性能最佳。特别是对于较高负载的参数空间，加速更为显着。在V100 GPU上，最有力的竞争对手是TensorRT，但在20个案例中，有13个案例比TensorRT更好。在工作量较大的情况下，TurboTransformer的优势尤其明显。

图6，固定长度请求上的三个模型的运行时的延迟。

批处理算法对吞吐量的提升效果如图6所示。图6的x轴表示每秒有多少请求到达服务系统，范围为20 req / sec至1500 req / sec。图12的y轴表示每秒可以获取多少响应。服务延迟保持稳定的关键点是请求吞吐量和响应吞吐量相等。当请求吞吐量高于临界点时，请求将累积在消息队列中，这将导致后面的请求长时间延迟。我们观察到，即使较差版本的实现（PyTorch-NoBatch）的吞吐量明显比TF服务高效，后者必须将序列填充到最大长度并带来巨大的开销。论文提出的批处理策略Turbo-DP-Batch在传统策略Turbo-Naive-Batch失效的情况下，仍然带来非常可观的吞吐量提升。

图7，不同请求吞吐量下服务的响应吞吐量。

论文最后分析了新的批处理算法对延迟的影响。如表1所示，在相同请求吞吐量下，新的批处理算法会带来更低的延迟。

表1，四个服务系统的延迟。表格中的每个项目代表平均（最小，最大）延迟（以毫秒为单位）。

总 结

TurboTransformers改进了在GPU数据中心中部署Tansformer服务的延迟和吞吐量。它解决了两个关键问题，即Transformer带来的巨大计算压力和可变长度输入的不确定性。为此，论文出了计算、内存和服务三个层次的优化方法。作为一个运行时库，TurboTransformers比PyTorch、TensorFlow-XLA、onnxruntime和FasterTransformers有着更加的性能表现。对实时的BERT线上服务，TurboTransformers也显著优于TF-serving等传统实现方案。

未来工作

未来，TurboTransformer希望探索和蒸馏、量化等算法优化结合，打造端到端的算法工程优化方案。

* PPoPP是并行计算领域的最重要会议之一，并行计算和编译技术领域很多经典的工作都是首先在 PPoPP 上发表的，它也是中国计算机学会 (CCF)推荐的 A 类国际学术会议，每年录取文章平均不超过30篇，录取率在20%左右。会议将在2021年 2月27到3月3期间在线上召开。