TensorFlow+TVM优化NMT神经机器翻译

TensorFlow+TVM优化NMT神经机器翻译
背景
神经机器翻译(NMT)是一种自动化的端到端方法,具有克服传统基于短语的翻译系统中的弱点的潜力。本文为全球电子商务部署NMT服务。
目前,将Transformer用作NMT系统的主要骨干,对基于经典RNN / LSTM模型的同等(甚至更高)精度进行高效的离线训练更为友好。尽管Transformer在离线训练阶段很友好,打破了跨时间步长的依赖性,但在线推理效率不高。在生产环境中,已经发现,初始版本的Transformer的推理速度约为1.5倍至2倍比LSTM版本慢。进行一些优化来提高推理性能,例如图形级op融合,循环不变节点运动。一个特殊挑战是,批处理matmul是Transformer中的主要性能热点,cuBLAS中的实现并未得到很好的优化。
在这里插入图片描述

下面的结果表明TVM生成内核(与调度表优化)带来至少13X加速批量MATMUL计算和futher加快与运营商的融合功能。
在这里插入图片描述

批处理Matmul
为什么批量matmul
在Transformer中,批处理matmul广泛用于计算多头注意。使用批处理matmul,层中的多个头部并行运行,帮助提高硬件的计算效率。
在这里插入图片描述

在推理阶段对Transformer模型进行了全面的分析,结果表明,批量matmul计算贡献了约30%的GPU内核执行时间。使用nvprof [2]对cuBLAS批处理matmul内核进行一些第一性原理分析,清楚地表明,当前的实现方式表现不佳,并且观察到了一些有趣的现象。
什么是批量matmul
批量矩阵计算对一批矩阵执行矩阵矩阵乘法。批处理被认为是“统一的”,即所有实例的A,B和C矩阵具有相同的大小(M,N,K),前导大小(Ida,Idb,Idc)和置换。
批量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)
}
}
批处理matmulshanpe
在语言翻译任务中,批处理matmul的shanpe比其它工作负载中的常规matmul计算小得多。Transformer中的shanpe与输入语句的长度和解码器步长有关。通常小于30。
对于批处理大小,给定一定的推理批处理大小,它是一个固定的数字。例如,如果将16用作光束大小为4的批大小,则批大小为16 * 4 * #head(多头注意中的头数,通常为8)。矩阵M,K,N的shanpe在[1,最大解码长度]或[1,最大编码长度]的范围内。
cuBLAS批次batch matmul的性能问题
首先,对批处理matmul内核进行了理论上的FLOP分析。结果非常有趣:所有批处理批量具有有限的计算强度(少于1个TFLOP)。
然后,通过nvprof剖析了具有多种shanpe的批处理matmul的cuBLAS性能。下表显示了在带有CUDA8.0的NVIDIA M40 GPU上获得的一些指标。
在这里插入图片描述

即使具有不同的shanpe(M,N,K不同),所有maxwell_sgemmBatched_128x128_raggedMn_tn调用也执行相同数量的FLOP,这比理论值大得多。所有这些不同的shanpe都可以被填充成一定的shanpe。在所有这些shanpe中,即使在最佳情况下,理论上的FLOP仍仅是实际执行的FLOP的2.74%,大多数计算是相当多余的。另一个cuBLAS内核maxwell_sgemmBatched_64x64_raggedMn_tn的调用也显示相同的现象。
cuBLAS的批量matmul实施远非效率。因此,使用TVM为NMT工作负载生成有效的批处理matmul内核。
批量matmul计算
在TVM中,一般的批量Matmul计算可以声明为:

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)
融合批处理mulmul的外部大小,即op大小的BB和FF,在批处理matmul计算中通常称为“批处理”大小。然后将外部和内部大小除以(number_thread * vthread)。
在批处理matmul中不需要交错模式,因此虚拟线程号(vthread_y和vthread_x)都设置为1。
寻找number_thread的最佳组合
以下结果是在具有CUDA8.0的NVIDIA M40 GPU设备上获得的。
在这里插入图片描述

找到的最佳组合num_thread_y和num_thread_x,通过强力搜索。经过强力搜索后,可以找到当前shanpe的最佳组合,该组合为num_thread_y= 8和num_thread_x= 32。
将batch matmul与其它算子融合在一起
现有的“黑盒” cuBLAS库调用充当通常使用的“ op Fusion”优化策略的边界。但是,利用所生成的高效批处理matmul核,可以容易地打破融合边界,不仅可以融合元素算子,还可以进一步提高性能。
从计算图可以看出,批处理matmul总是跟在广播添加算子或转置算子之后。通过将“ add”或“ transpose”算子与批处理mult融合,可以减少内核启动开销和冗余内存访问时间。
批处理matmul和广播添加融合计算可以声明如下:

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)
批处理matmul和转置融合计算可以声明为:

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]的shanpe以详细说明所生成代码的性能。选择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在工作量中的输入shanpe是有限的,可以很容易地预先枚举。使用这些预定义的shanpe,可以提前生成高度优化的CUDA内核(固定shanpe计算可以带来最佳的优化潜力)。同时,还将生成适用于大多数shanpe的通用批处理matmul内核,为没有相应的提前生成的内核的shanpe提供回退机制。
针对特定shanpe生成的高效内核和后备内核已集成到Tensorflow框架中。开发融合算子,例如BatchMatMulTranspose或BatchMatMulAdd,以使用TVM的运行时API针对特定输入shanpe启动特定生成的内核,或调用后备内核。进行图形优化遍历,用融合的算子自动替换原始批处理matmul +添加/转置模式。结合更积极的图形优化过程,试图利用TVM为长尾算子模式生成更有效的融合内核,以进一步提高端到端性能。
总结
在阿里巴巴内部,发现TVM是开发高性能GPU内核以满足内部需求的非常有效的工具。以NMT变压器模型为例来说明使用TVM的优化策略。首先,通过第一性原理分析确定了Transformer模型的热点。然后使用TVM生成高度优化的CUDA内核来取代CUBLAS版本(13X加速观察)。接下来,利用TVM的内核融合机制融合批处理matmul的先前/以下算子,以进一步提高性能(进一步提高1.7倍的性能)。端到端性能提高了1.4倍。基于这些生成的内核,开发了图优化遍历,自动用TVM融合内核替换原始计算模式,确保优化对最终用户是透明的,因为作为AI基础架构提供商,发现优化策略的透明性对于推广其优化算法非常重要。最后,但并非最不重要的一点是,所有这些优化都以松散耦合的方式集成到TensorFlow中,展示了将TVM与不同的深度学习框架集成的潜在方法。此外,将TVM集成为TensorFlow的代码源后端的工作,希望将来能共享更多结果。

已标记关键词 清除标记
相关推荐
©️2020 CSDN 皮肤主题: 深蓝海洋 设计师:CSDN官方博客 返回首页