阿里巴巴最新实践：TVM+TensorFlow提高神经机器翻译性能
更多干货内容请关注微信公众号“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 倍，并且在启用了算子融合的情况下还将进一步提高速度。
在 Transformer 中，批量 matmul 被广泛应用于多头 attention（multi-head attention）的计算。利用批量 matmul，attention 层中的多头可以并行运行，有助于提高硬件的计算效率。
我们在推理阶段对 Transformer 模型进行了深入的剖析，结果表明批量 matmul 计算占 GPU 内核执行时间的 30% 左右。我们利用 nvprof[2] 对 cuBLAS 的批量 matmul 内核进行了第一性原理分析，结果表明当前的实现效果很差，同时我们还发现了一些有趣的现象。
通常，批量 matmul 计算是指在一批矩阵上执行矩阵 - 矩阵乘法。其中 batch 被认为是“均匀的”，即所有实例对于它们各自的 A、B 和 C 矩阵具有相同的维度（M、N、K）、前导维度（lda、ldb、ldc）和转换率。
批量 matmul 计算可以更具体地描述如下：
在语言翻译任务中，批量 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, 最大编码长度] 的范围内。
首先，我们对批量 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 内核。
在 TVM 中，通常批量 matmul 计算可以声明为：
在对批量 matmul 计算进行声明之后，我们需要仔细设计Schedule，以充分挖掘其性能潜力。
我们融合了批量 matmul 的外部维度，即 op 维度的 BB 和 FF，在批量 matmul 计算中通常称为“批量”维度。然后我们用因子（number _ thread * vthread）分割外部维度和内部维度。
批量 matmul 中不需要 stridged 模式，因此虚拟线程数（vthread_y 和 vthread_x）都设置为 1。
下图是使用 CUDA8.0 在 NVIDIA M40 GPU 上得到的结果。
根据以往的经验，寻找 num_thread_y 和 num_thread_x 最佳组合的方法是暴力搜索。通过暴力搜索，可以找到当前形状（shape）的最佳组合，我们在当前计算中得到的最佳组合为 num_thread_y=8，num_thread_x=32。
当前的“黑盒”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 倍。
在我们的工作负载中，批量 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
今日荐文
点击下方图片即可阅读
IBM全新机器学习库 SnapML：训练提速 46倍，40亿样本仅需 91.5秒
活动推荐
TGO 鲲鹏会是极客邦科技 / InfoQ 旗下面向高端技术管理者的微信公众号，每日分享最全面的技术管理干货与技术圈一线大咖分享。发布过近千篇国内外知名 CTO 独家干货分享、经验总结、专访等高质量原创内容，致力于帮助技术人持续学习成长。
Tips：现在关注「 TGO鲲鹏会 」微信公众号即可获得 CTO 技能图谱一份。
一些粉丝朋友还没有养成阅读后点赞的习惯，希望大家在阅读后顺便点赞，以示鼓励！原创是一种信仰，专注是一种态度！