TVM优化GPU机器翻译
TVM优化GPU机器翻译
背景
神经机器翻译（NMT）是一种自动化的端到端方法，具有克服传统基于短语的翻译系统中的弱点的潜力。最近，阿里巴巴集团正在为全球电子商务部署NMT服务。
将Transformer用作NMT系统的关键技术，相对于基于经典RNN / LSTM的模型具有同等（甚至更高）的精度，对于高效的离线训练更为友好。尽管Transformer在离线训练阶段很友好，打破了跨时间步长的依赖性，但在线推理效率不高。在生产环境中，已经发现，初始版本的Transformer的推理速度约为1.5倍至2倍，比LSTM版本慢。已经进行了一些优化来提高推理性能，例如图形级op融合，循环不变节点运动。观察到的一个特殊挑战是，批处理matmul是Transformer中的主要性能热点，而cuBLAS中的当前实现并未得到很好的优化。
下面的结果表明TVM生成内核（与调度表优化）带来至少13X批量MATMUL计算加速，futher加快算子融合功能。
批处理Matmul
为什么批量Matmul
在Transformer中，批处理matmul广泛用于multi-head attention计算。使用matmul批处理，关注层中multiple heads并行运行，帮助提高硬件的计算效率。
在推理阶段对Transformer模型进行了全面的分析，结果表明，批量matmul计算贡献了约30％的GPU内核执行时间。使用nvprof，对cuBLAS批处理matmul内核进行一些第一性原理分析，可以清楚地表明，当前的实现方式表现不佳，并且观察到了一些有趣的现象。
什么是批量matmul
通常，批量矩阵计算对一批矩阵执行乘法。批处理被认为是“统一的”，即所有实例具有相同的维度（M，N，K），前导维度（lda，ldb，ldc），以及各自的A，B和C矩阵的转置。
批量matmul计算，可以更具体地描述如下：
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) }}批处理matmul形状
在语言翻译任务中，在其它工作负载中，批处理matmul的形状显着小于常规matmul计算。Transformer中的形状与输入语句的长度和解码器步长有关。通常小于30。
对于批处理尺寸，给定一定的推断批处理大小，这是一个固定的数字。例如，如果将16用作光束大小为4的批大小，则批大小为16 * 4 * #head（多头注意中的头数，通常为8）。矩阵M，K，N的形状，在[1，最大解码长度]或[1，最大编码长度]的范围内。
cuBLAS batch matmul的性能问题
首先，对批处理matmul内核进行了理论上的FLOP分析。结果非常有趣：所有批处理Matmul的计算强度都有限（少于1个TFLOP）。
然后，通过nvprof剖析了具有多种形状的matmul批处理cuBLAS性能。下表显示了在带有CUDA8.0的NVIDIA M40 GPU上获得的一些指标。
|
输入形状 |
核心 |
理论FLOP |
nvprof观察到的FLOP |
理论FLOPs / |
|
[512、17、17、128] |
maxwell_sgemmBatched_128x128_raggedMn_tn |
18939904 |
2155872256 |
0.87％ |
|
[512、1、17、128] |
maxwell_sgemmBatched_128x128_raggedMn_tn |
1114112 |
2155872256 |
0.052％ |
|
[512，17，1，128] |
maxwell_sgemmBatched_128x128_raggedMn_tn |
1114112 |
2155872256 |
0.052％ |
|
[512，30，30，128] |
maxwell_sgemmBatched_128x128_raggedMn_tn |
58982400 |
2155872256 |
2.74％ |
即使具有不同的形状（M，N，K不同），所有maxwell_sgemmBatched_128x128_raggedMn_tn调用，也执行相同数量的FLOP，这比理论值大得多。可以推断出，所有这些不同的形状都可以被填充成一定的形状。在所有这些形状中，即使在最佳情况下，理论上的FLOP仍然仅是实际执行的FLOP的2.74％，因此，大多数计算都是相当多余的。同样，另一个cuBLAS内核maxwell_sgemmBatched_64x64_raggedMn_tn的调用也显示相同的现象。
显然，cuBLAS的批量matmul实施远非效率。因此，使用TVM为NMT工作负载生成有效的批处理matmul内核。
批量matmul计算
在TVM中，一般的批量Matmul计算可以声明为：
# computation representationA = 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)融合了批处理matmul的外部尺寸，即op尺寸的BB和FF，在批处理matmul计算中通常称为“批处理”尺寸。然后将外部和内部尺寸除以（number_thread * vthread）。
在批处理matmul中不需要交错模式，因此虚拟线程号（vthread_y和vthread_x）都设置为1。
寻找number_thread的最佳组合
以下结果是在具有CUDA8.0的NVIDIA M40 GPU设备上获得的。
|
输入形状[批处理，特征，M，N，K] |
num_thread_y，num_thread_x |
num_vthread_y，num_vthread_x |
时间（us） |
|
[64,8,1,17,128] |
8,1 |
32,1 |
37.62 |
|
[64,8,1,17,128] |
4,1 |
32,1 |
39.30 |
|
[64,8,1,17,128] |
1,1 |
32,1 |
38.82 |
|
[64,8,1,17,128] |
1,1 |
256,1 |
41.95 |
|
[64,8,1,17,128] |
32,1 |
1,1 |
94.61 |
从以往的经验了解到，该方法找到的最佳组合num_thread_y，num_thread_x通过强力搜索。经过蛮力搜索后，可以找到当前形状的最佳组合，在当前计算中为num_thread_y= 8和num_thread_x= 32。
将批处理matmul 与其它算子融合
通常，现有的“黑盒” cuBLAS库，调用扮演着通常使用的“ op Fusion”优化策略的边界的角色。但是，利用所生成的高效批处理matmul核，可以容易地打破融合边界，不仅可以融合元素的操作方式，还可以进一步提高性能。
从计算图可以看出，批处理matmul总是跟在添加操作或转置操作广播之后。通过将“ add”或“ transpose”操作与批处理matmul融合，可以减少内核启动开销和冗余内存访问时间。
批处理matmul和添加融合广播计算，可以声明如下：
# computation representationA = 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 patternB = 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)批处理matmul和转置融合计算可以声明为：
# computation representationA = 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 - TF-R1.4
BatchMatmul+Transpose（另购）：541.9 - TVM
BatchMatmul：37.62美元 - TVM
BatchMatmul+Transpose（融合）：38.39美元
内核融合优化进一步提高了1.7倍的速度。
与Tensorflow集成
批量matmul在工作量中的输入形状是有限的，可以很容易地预先枚举。使用这些预定义的形状，可以提前生成高度优化的CUDA内核（固定形状计算可以带来最佳的优化潜力）。将生成适用于大多数形状的通用批处理matmul内核，为没有相应的提前生成的内核的形状提供回退机制。
针对特定形状生成的高效内核和后备内核已集成到Tensorflow框架中。开发了融合操作，例如BatchMatMulTranspose或BatchMatMulAdd，使用TVM的runtime API生成特定生成的内核，实现某些输入形状或调用后备内核。进行图形优化遍历，用算子融合自动替换原始批处理matmul +添加/转置模式。通过结合更积极的图形优化过程，试图利用TVM为长尾算子模式生成更有效的融合内核，以进一步提高端到端性能。
概括
在阿里巴巴内部，发现TVM是开发高性能GPU内核，满足内部需求的非常有效的工具。以NMT Transformer模型为例来说明使用TVM的优化策略。首先，通过第一性原理分析确定了Transformer模型的热点。然后使用TVM生成高度优化的CUDA内核来取代CUBLAS版本（13X加速观察）。接下来，利用TVM的内核融合机制融合批处理matmul的先前/以下操作，以进一步提高性能（进一步提高1.7倍的性能）。端到端性能提高了1.4倍。基于这些生成的内核，开发了图优化遍历以自动用TVM融合内核替换原始计算模式，确保优化对最终用户是透明的，作为AI基础设施提供商，发现优化策略的透明性对于推广其优化算法非常重要采用。最后，并非最不重要的一点是，所有这些优化都以松散耦合的方式集成到TensorFlow中，展示了将TVM与不同的深度学习框架集成的潜在方法。此外，正在进行一项将TVM集成为TensorFlow的代码源后端的工作，希望将来能与社区共享更多结果。
