阿里巴巴最新实践：TVM+TensorFlow提升神经机器翻译性能

策划 & 编译｜Natalie

AI 前线导读： 本文是阿里巴巴 PAI-Blade 团队发表于 TVM 的最新博文，文中阐述了如何将 TVM 引入 TensorFlow，使 TensorFlow 中的 batchmul 速度提升 13 倍，同时将端到端神经机器翻译性能提升 1.7 倍。AI 前线对原文进行了编译。

更多干货内容请关注微信公众号“AI 前线”，（ID：ai-front）

TVM 是由华盛顿大学在读博士陈天奇等人提出的深度学习自动代码生成方法，该技术能自动为大多数计算硬件生成可布署优化代码，其性能可与当前最优的供应商提供的优化计算库相比，且能够适应新型专用加速器后端。

背景

神经机器翻译（NMT）是一种端到端的自动翻译方法，有可能克服传统短语翻译系统的不足。最近，阿里巴巴集团正在尝试将 NMT 服务部署到全球电子商务业务场景中。

目前，咱们正在尝试开发 Transformer[1] 做为咱们 NMT 系统的核心组件，由于它比基于 RNN/LSTM 的经典模型更有利于高效的离线训练，具备同等（甚至更高）的精确度。虽然 Transformer 打破了时间步长之间的依赖关系，所以对于离线训练阶段很是友好，可是对于在线推理来讲，就没有那么高效了。在咱们的生产环境中，初始版本的 Transformer 推理速度比 LSTM 版本慢大约 1.5 倍到 2 倍。为了提升推理性能，咱们已经尝试进行了多种优化，如图级别 op 融合、循环不变式节点运动 [3]。咱们发现批量 matmul 是 Transformer 中的一个性能关键点，而当前 cuBLAS 的实现并未进行很好的优化。

咱们的实践结果代表，经过 TVM 生成的内核（优化了 schedule ）可使批量 matmul 计算速度提升至少 13 倍，而且在启用了算子融合的状况下还将进一步提升速度。

批量 Matmul

为何要使用批量 matmul

在 Transformer 中，批量 matmul 被普遍应用于多头 attention（multi-head attention）的计算。利用批量 matmul，attention 层中的多头能够并行运行，有助于提升硬件的计算效率。

咱们在推理阶段对 Transformer 模型进行了深刻的剖析，结果代表批量 matmul 计算占 GPU 内核执行时间的 30% 左右。咱们利用 nvprof[2] 对 cuBLAS 的批量 matmul 内核进行了第一性原理分析，结果代表当前的实现效果不好，同时咱们还发现了一些有趣的现象。

什么是批量 matmul

一般，批量 matmul 计算是指在一批矩阵上执行矩阵 - 矩阵乘法。其中 batch 被认为是“均匀的”，即全部实例对于它们各自的 A、B 和 C 矩阵具备相同的维度（M、N、K）、前导维度（lda、ldb、ldc）和转换率。

批量 matmul 计算能够更具体地描述以下：

批量 matmul 的 shape

在语言翻译任务中，批量 matmul 的 shape 明显小于其余工做负载下的正常 matmul 计算。Transformer 的 shape 与输入句子的长度和解码器步骤的长度有关。一般会小于 30。

当给定推理的 batch size 时，批量维度（batch dimension）是一个固定值。例如，若是将 batch size 设定为 16，beam 大小为 4，则批量维度为 16*4*#head（多头 attention 中的头数，一般为 8），矩阵 M、K、N 的 shape 在 [1, 最大解码长度] 或 [1, 最大编码长度] 的范围内。

cuBLAS 批量 matmul 的性能问题

首先，咱们对批量 matmul 内核进行了 FLOPs 理论分析。结果至关有趣：全部批量 matmul 的计算强度都很是有限（小于 1TFLOPs）。

接着咱们利用 nvprof 软件对不一样 shape 的批量 matmul 的 cuBLAS 进行了分析。下表显示了使用 CUDA 8.0 在 NVIDIA M40 GPU 上获得的一些指标。

即便使用不一样的 shape（改变 M , N , K），全部的 Maxwell_sgembatched_128x128_raggedMn_tn 调用都会执行相同数量的 FLOP，并且实际值比理论值大得多。能够推断，全部这些不一样的 shape 可能都被填充成了特定的 shape。在全部这些 shape 中，即便在最好的状况下，理论 FLOP 仍仅占实际执行 FLOP 的 2.74%，所以大部分计算至关冗余。相似地，调用另外一个 cuBLAS 内核的 Maxwell_sgembatched_64x64_raggedMn_tn 出现了一样的现象。

显而易见，cuBLAS 的批量 matmul 实现效率很低。所以，咱们使用 TVM 为 NMT 工做负载生成高效的批量 matmul 内核。

批量 matmul 计算

在 TVM 中，一般批量 matmul 计算能够声明为：

Schedule优化

在对批量 matmul 计算进行声明以后，咱们须要仔细设计Schedule，以充分挖掘其性能潜力。

调参：块 / 线程数

咱们融合了批量 matmul 的外部维度，即 op 维度的 BB 和 FF，在批量 matmul 计算中一般称为“批量”维度。而后咱们用因子（number _ thread * vthread）分割外部维度和内部维度。

批量 matmul 中不须要 stridged 模式，所以虚拟线程数（vthread_y 和 vthread_x）都设置为 1。

寻找 number_thread 的最佳组合

下图是使用 CUDA8.0 在 NVIDIA M40 GPU 上获得的结果。

根据以往的经验，寻找 num_thread_y 和 num_thread_x 最佳组合的方法是暴力搜索。经过暴力搜索，能够找到当前形状（shape）的最佳组合，咱们在当前计算中获得的最佳组合为 num_thread_y=8，num_thread_x=32。

将批量 matmul 与其余运算融合

当前的“黑盒”cuBLAS 库调用一般会被做为经常使用“op 融合”优化策略的边界。然而，利用生成的高效批量 matmul 内核，能够很容易地打破这一融合边界，不只能够融合 element-wise 运算，从而进一步提升性能。

从计算图中能够看出，批量 matmul 以后老是会进行广播加法运算或转置运算。经过将“加法”或“转置”运算与批量 matmul 融合，能够减小内核启动开销和冗余内存访问时间。

批量 matmul 和广播加法融合计算可声明以下：

批量 matmul 和转置融合计算可声明以下：

融合后的内核性能

咱们选择 [batch=64, heads=8, M=1, N=17, K=128] 的形状来测试前文生成的代码的性能。咱们将序列长度设置为 17，由于这个值是咱们生产场景中的平均输入长度。

测试结果以下：

• tf-r1.4 BatchMatmul：513.9 us

• tf-r1.4 BatchMatmul+Transpose (separate)：541.9 us

• TVM BatchMatmul：37.62 us

• TVM BatchMatmul+Transpose (fused)：38.39 us

内核融合优化进一步将速度提升了 1.7 倍。

与 TensorFlow 集成

在咱们的工做负载中，批量 matmul 的输入形状是有限的，而且易于预先枚举。利用这些预约义的形状，咱们能够提早生成高度优化的 CUDA 内核（固定形状计算能够带来最佳的优化潜力）。同时，咱们还将生成适用于大多数形状的通用批量 matmul 内核，为没有提早生成内核的形状提供回退机制。

咱们将生成的针对特定形状的高效内核和回退机制集成到了 TensorFlow 框架中。咱们开发了融合op，如 BatchMatMulTranspose 或 BatchMatMulAdd，以使用 TVM 的运行时 API 为特定输入形状启动特定的生成内核，或调用回退内核。经过执行图形优化通道，能够用融合 op 自动替换原始的批量 matmul+ 加法 / 转置模式。同时，经过结合更激进的图形优化通道，咱们尝试利用 TVM 为长尾操做模式生成更高效的融合内核，以进一步提升端到端性能。

总结

在阿里巴巴内部，咱们发现 TVM 是一个很是高效的工具，咱们能够用它开发高性能的 GPU 内核以知足内部需求。

本文以 NMT Transformer 模型为例，对 TVM 的优化策略进行了详细说明。首先，经过第一性原理分析找出 Transformer 模型的关键问题。而后使用 TVM 生成高度优化的 CUDA 内核来替换 cuBLAS 版本（此时已经获得了 13 倍的加速效果）。接下来，咱们利用 TVM 的内核融合机制来融合批量 matmul 的前 / 后操做，从而带来进一步的性能改善（性能进一步提升 1.7 倍）。在此基础上，咱们开发了一个图形优化通道，自动用 TVM 融合内核替换原有的计算模式，保证优化过程对最终用户透明。做为 AI 基础设施提供商，咱们发现透明度对于推广优化策略的应用很是重要。

最后，全部这些优化都以松散耦合的方式集成到了 TensorFlow 中，展现了将 TVM 与不一样深度学习框架集成的一种可能方式。此外，咱们目前还在进行将 TVM 整合为 TensorFlow 代码生成后端的工做，但愿从此可以与社区分享更多成果。

资源

TVM implementation of fused batch matmul + transpose computation项目代码

https://github.com/Orion34C/tvm-batch-matmul-example/blob/master/tvm_batch_matmul_transpose_m1_kX.py

PAI 团队中文介绍

https://zhuanlan.zhihu.com/p/33513874

参考文献

[1] Attention is All You Need

https://arxiv.org/pdf/1706.03762.pdf

[2] nvprof is Your Handy Universal GPU Profiler

https://devblogs.nvidia.com/cuda-pro-tip-nvprof-your-handy-universal-gpu-profiler/

[3] Add Loop Invariant Node Motion Optimization in GraphOptimizer

https://github.com/tensorflow/tensorflow/pull/16306

原文连接：

http://www.tvmlang.org/2018/03/23/nmt-transformer-optimize.html

更多干货内容请关注微信公众号“AI 前线”，（ID：ai-front）