将 TVM 引入 TensorFlow 以优化 GPU 上的神经机器翻译

2018 年 3 月 23 日

作者

这是一篇由阿里巴巴集团机器翻译平台团队和 PAI-Blade 团队投稿的客座博客文章

背景

神经机器翻译 (NMT) 是一种用于自动化翻译的端到端方法，它有可能克服传统基于短语的翻译系统的弱点。最近，阿里巴巴集团正在致力于为全球电子商务部署 NMT 服务。

目前，我们正在采用 Transformer [1] 作为 NMT 系统中的主要骨干网络，因为它更利于高效的离线训练，并且精度与经典的基于 RNN/LSTM 的模型相当（甚至更高）。虽然 Transformer 对于离线训练阶段很友好，因为它打破了跨时间步的依赖关系，但对于在线推理而言效率不高。在我们的生产环境中，我们发现初始版本的 Transformer 的推理速度比 LSTM 版本慢约 1.5 倍到 2 倍。我们已经采取了一些优化措施来提高推理性能，例如图级算子融合、循环不变节点移动 [3]。我们观察到的一个特殊挑战是，批量矩阵乘法是 Transformer 中的主要性能热点，而 cuBLAS 中当前的实现并没有得到很好的优化。

以下结果表明，TVM 生成的内核（通过调度优化）为批量矩阵乘法计算带来了至少 13 倍 的加速，并且在启用算子融合后进一步加速。

批量矩阵乘法

为什么是批量矩阵乘法

在 Transformer 中，批量矩阵乘法广泛用于多头注意力机制的计算。使用批量矩阵乘法，注意力层中的多个头可以并行运行，这有助于提高硬件的计算效率。

我们对 Transformer 模型在推理阶段进行了全面的性能分析，结果表明，批量矩阵乘法计算占 GPU 内核执行时间的 ~ 30%。使用 nvprof[2] 对 cuBLAS 的批量矩阵乘法内核进行一些第一性原理分析，清楚地表明当前的实现性能不佳，并观察到一些有趣的现象。

什么是批量矩阵乘法

通常，批量矩阵乘法计算对一批矩阵执行矩阵-矩阵乘法。该批次被认为是“统一的”，即所有实例都具有相同的维度 (M, N, K)、前导维度 (lda, ldb, ldc) 以及各自 A、B 和 C 矩阵的转置。

批量矩阵乘法计算可以更具体地描述如下

void BatchedGemm(input A, input B, output C, M, N, K, batch_dimension) {
  for (int i = 0; i < batch_dimension; ++i)  {
    DoGemm(A[i],B[i],C[i],M,K,N)
  }
}

批量矩阵乘法形状

在语言翻译任务中，批量矩阵乘法的形状明显小于其他工作负载中正常的矩阵乘法计算。Transformer 中的形状与输入句子的长度和解码步骤的长度有关。通常，它小于 30。

至于批次维度，它是一个固定的数字，由给定的推理批次大小决定。例如，如果使用 16 作为批次大小，波束大小为 4，则批次维度为 16 * 4 * #head（多头注意力机制中的头数，通常为 8）。矩阵 M、K、N 的形状在 [1, 最大解码长度] 或 [1, 最大编码长度] 的范围内。

cuBLAS 批量矩阵乘法的性能问题

首先，我们对批量矩阵乘法内核进行了理论 FLOPs 分析。结果非常有趣：所有批量矩阵乘法的计算强度都有限（小于 1 TFLOPs）。

然后，我们通过 nvprof 分析了 cuBLAS 批量矩阵乘法在多种形状下的性能。下表显示了在 NVIDIA M40 GPU（CUDA 8.0）上获得的一些指标。

即使形状不同（M、N、K 各异），所有 maxwell_sgemmBatched_128x128_raggedMn_tn 调用都执行相同数量的 FLOPs，这远大于理论值。可以推断，所有这些不同的形状都可能被填充到某个特定形状。在所有这些形状中，即使在最佳情况下，理论 FLOPs 仍然只占实际执行的 FLOPs 的 2.74%，因此大部分计算都是非常冗余的。类似地，另一个 cuBLAS 内核 maxwell_sgemmBatched_64x64_raggedMn_tn 的调用也显示出相同的现象。

显然，cuBLAS 的批量矩阵乘法实现效率低下。因此，我们使用 TVM 为我们的 NMT 工作负载生成高效的批量矩阵乘法内核。

批量矩阵乘法计算

在 TVM 中，通用的批量矩阵乘法计算可以声明为

# computation representation
A = tvm.placeholder((batch, M, K), name='A')
B = tvm.placeholder((batch, K, N), name='B')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute((batch, M, N),
         lambda b, y, x: tvm.sum(A[b, y, k] * B[b, k, x], axis = k),
         name = 'C')

调度优化

在声明计算之后，我们需要仔细设计我们自己的调度，以挖掘性能潜力。

调整块/线程数的参数

  # thread indices
  block_y = tvm.thread_axis("blockIdx.y")
  block_x = tvm.thread_axis("blockIdx.x")
  thread_y = tvm.thread_axis((0, num_thread_y), "threadIdx.y")
  thread_x = tvm.thread_axis((0, num_thread_x), "threadIdx.x")
  thread_yz = tvm.thread_axis((0, vthread_y), "vthread", name="vy")
  thread_xz = tvm.thread_axis((0, vthread_x), "vthread", name="vx")

  # block partitioning
  BB, FF, MM, PP = s[C].op.axis
  BBFF = s[C].fuse(BB, FF)
  MMPP = s[C].fuse(MM, PP)
  by, ty_block = s[C].split(BBFF, factor = num_thread_y * vthread_y)
  bx, tx_block = s[C].split(MMPP, factor = num_thread_x * vthread_x)
  s[C].bind(by, block_y)
  s[C].bind(bx, block_x)
  vty, ty = s[C].split(ty_block, nparts = vthread_y)
  vtx, tx = s[C].split(tx_block, nparts = vthread_x)
  s[C].reorder(by, bx, vty, vtx, ty, tx)
  s[C].reorder(by, bx, ty, tx)
  s[C].bind(ty, thread_y)
  s[C].bind(tx, thread_x)
  s[C].bind(vty, thread_yz)
  s[C].bind(vtx, thread_xz)

我们融合了批量矩阵乘法的外部维度，即操作维度的 BB 和 FF，通常在批量矩阵乘法计算中被称为“批次”维度。然后，我们将外部维度和内部维度按 (number_thread * vthread) 的因子进行拆分。

批量矩阵乘法中不需要步幅模式，因此虚拟线程数（vthread_y 和 vthread_x）都设置为 1。

找到 number_thread 的最佳组合

以下结果是在 NVIDIA M40 GPU 设备上使用 CUDA 8.0 获得的。

从过去的经验中学习到，找到 num_thread_y 和 num_thread_x 的最佳组合的方法是通过暴力搜索。经过暴力搜索后，可以找到当前形状的最佳组合，在当前计算中为 num_thread_y = 8 和 num_thread_x = 32。

将批量矩阵乘法与其他操作融合

通常，现有的“黑盒” cuBLAS 库调用充当通常使用的“算子融合”优化策略的边界。然而，借助生成的有效批量矩阵乘法内核，可以轻松打破融合边界，可以融合的不仅仅是逐元素操作，从而可以获得进一步的性能提升。

从计算图中观察到，批量矩阵乘法之后总是跟随一个广播加法操作或一个转置操作。通过将“加法”或“转置”操作与批量矩阵乘法融合，可以减少内核启动开销和冗余内存访问时间。

批量矩阵乘法和广播加法融合计算可以声明如下

# computation representation
A = tvm.placeholder((batch_size, features, M, K), name='A')
# the shape of B is (N, K) other than (K, N) is because B is transposed is this fusion pattern
B = tvm.placeholder((batch_size, features, N, K), name='B')
ENTER = tvm.placeholder((batch_size, 1, M, N), name = 'ENTER')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute(
           (batch_size, features, M, N),
           lambda yb, yf, m, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, x, k], axis = k),
           name = 'C')
D = topi.broadcast_add(C, ENTER)

批量矩阵乘法和转置融合计算可以声明为

# computation representation
A = tvm.placeholder((batch_size, features, M, K), name='A')
B = tvm.placeholder((batch_size, features, K, N), name='B')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute(
           (batch_size, M, features, N),
           lambda yb, m, yf, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, k, x], axis = k),
           name = 'C')

融合内核性能

选择形状为 [batch=64, heads=8, M=1, N=17, K=128] 来详细说明生成代码的性能。选择 17 作为序列长度，因为它是我们生产场景中的平均输入长度。

• tf-r1.4 BatchMatmul: 513.9 us
• tf-r1.4 BatchMatmul + Transpose (分离): 541.9 us
• TVM BatchMatmul: 37.62 us
• TVM BatchMatmul + Transpose (融合): 38.39 us

内核融合优化带来了进一步 1.7 倍 的加速。

与 Tensorflow 集成

我们的工作负载中批量矩阵乘法的输入形状是有限的，并且可以提前轻松枚举。借助这些预定义的形状，我们可以提前生成高度优化的 CUDA 内核（固定形状计算可以带来最佳的优化潜力）。同时，还将生成适用于大多数形状的通用批量矩阵乘法内核，为没有相应提前生成内核的形状提供回退机制。

针对特定形状生成的有效内核和回退内核已集成到 Tensorflow 框架中。我们开发了融合算子，例如 BatchMatMulTranspose 或 BatchMatMulAdd，以使用 TVM 的运行时 API 为特定输入形状启动特定的生成内核，或调用回退内核。执行图优化过程以自动将原始批量矩阵乘法 + 加法/转置模式替换为融合算子。同时，通过结合更积极的图优化过程，我们正在尝试利用 TVM 为长尾操作模式生成更高效的融合内核，以进一步加速端到端性能。

总结

在阿里巴巴内部，我们发现 TVM 是一个非常高效的工具，可以开发高性能 GPU 内核以满足我们的内部需求。在本博客中，以 NMT Transformer 模型为例来说明我们使用 TVM 的优化策略。首先，我们通过第一性原理分析定位 Transformer 模型的热点。然后，我们使用 TVM 生成高度优化的 CUDA 内核来替换 cuBLAS 版本（观察到 13 倍 的加速）。接下来，我们利用 TVM 的内核融合机制来融合批量矩阵乘法的前置/后置操作，以带来进一步的性能提升（进一步 1.7 倍 的性能提升）。端到端性能提升为 1.4 倍。基于这些生成的内核，我们开发了一个图优化过程，以自动将原始计算模式替换为 TVM 融合内核，以确保优化对最终用户是透明的，因为作为 AI 基础设施提供商，我们发现优化策略的透明度对于推广其应用非常重要。最后但并非最不重要的一点是，所有这些优化都以松耦合的方式集成到 TensorFlow 中，展示了将 TVM 与不同深度学习框架集成的潜在方法。此外，目前正在进行将 TVM 集成为 TensorFlow 的代码生成后端的 工作，我们希望未来能够与社区分享更多成果。

资源

• 融合批量矩阵乘法 + 转置计算的 TVM 实现

参考文献

[1] Attention is All You Need

[2] nvprof is Your Handy Universal GPU Profiler

[3] Add Loop Invariant Node Motion Optimization in GraphOptimizer