阿里将 TVM 融入 TensorFlow,在 GPU 上实现全面提速

AI 研习社按,日前,阿里机器翻译团队和 PAI 团队发表博文,阐述将 TVM 引入 TensorFlow,可以带来至少 13 倍的 batch 矩阵相乘(matmul)加速。雷锋网 AI 研习社将原文编译整理如下:

背景

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

目前,我们将 Transformer(https://arxiv.org/pdf/1706.03762.pdf) 作为 NMT 系统的核心组成。相较于传统基于 RNN/LSTM 的方法,它更适合于高效的离线训练,有着相同或更高的精度。

Transformer 在时间步长中打破了相关性,对离线训练更友好,但在在线推理上,它并没有那么高效。我们在生产环境中发现初版 Transformer 的推理速度大约比 LSTM 版本慢 1.5 倍到 2 倍。为了提高推理性能,我们已经进行了一些优化,包括图级别的 op 融合、循环不变节点外提(loop invariant node motion)。我们观察到一个特殊问题:batch 矩阵相乘是 Transformer 中的一个关键问题,目前它在 cuBLAS 中的实现并未得到很好的优化。

图1:Transformer 模型架构

下图表明,通过 TVM 生成的内核可以带来至少 13 倍的 batch 矩阵相乘加速,伴随算子融合,速度将更快。

batch 矩阵相乘

为什么选择利用 batch 矩阵相乘

在 Transformer 中,batch 矩阵相乘被广泛应用于 multi-head attention 的计算。利用 batch 矩阵相乘,可以并行运行 attention 层中的 multiple heads,这有助于提高硬件的计算效率。

图2:左图为 Scaled Dot-Product Attention,右图为并行运行若干 attention 层的 Multi-Head Attention

我们在推理阶段对 Transformer 模型进行了全面分析,结果表明,batch 矩阵相乘计算的开销达到 GPU 内核执行时间的 30%。当使用 nvprof 对 cuBLAS batch 矩阵相乘内核做一些第一原理(first-principle)分析,很明显,这种方法的表现并不好,同时我们还发现几个有趣的现象。

什么是 batch 矩阵相乘

通常,batch 矩阵相乘计算会在一批矩阵上执行矩阵-矩阵乘法。batch 被认为是「统一的」,即所有实例都具有相同的维度(M,N,K)、leading 维度 (lda,ldb,ldc) 和它们各自的 A、B、C 矩阵的转置。

batch 矩阵相乘计算具体可以描述如下:

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) } }

batch 矩阵相乘形状

在语言翻译任务中,batch 矩阵相乘的形状比在其他工作负载下的常规矩阵相乘计算要小得多。Transformer 的形状与输入语句的长度和解码器步长有关。一般来说小于 30。

至于 batch 维度,当给定推理 batch 大小时,它是固定数字。例如,如果 batch size 是 16,beam size 是 4,batch 维度是 16 * 4 * #head (在 multi-headattention 中 head 的数目,通常为 8)。矩阵 M、K、N 的范围在 [1, max decode length] 或 [1, max encode length] 内。

batch 矩阵相乘的性能问题

首先,我们在理论上对 batch 矩阵相乘内核进行了 FLOP 分析。结果非常有趣:所有 batch 矩阵相乘的计算强度都是受限的(TFLOP 数少于 1)。

然后,我们通过 nvprof 描述了多形状 batch 矩阵相乘的 cuBLAS 性能。下面的表格中是使用 NVIDIA M40 GPU(CUDA 8.0)得到的一些指标。

即使形状不同(在 M、N、K 间变化),所有 maxwell_sgemmBatched_128x128_raggedMn_tn 调用执行的都是相同的 FLOP 数,这比理论值大得多。从中可以推断,所有这些不同的形状最终都会被填充成确定的形状。在所有的形状中,即使在最好的情况下,理论 FLOP 只占实际执行 FLOP 的 2.74%,因此大多数计算都是多余的。类似地,调用另一个 cuBLAS 内核 maxwell_sgemmBatched_64x64_raggedMn_tn 也出现相同情况。

显而易见,cuBLAS batch 矩阵相乘的执行效率很低。基于这个原因,我们在 NMT 中使用 TVM 生成高效的 batch 矩阵相乘内核。

batch 矩阵相乘计算

在 TVM 中,普通 batch 矩阵相乘计算声明如下:

# 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')

调度优化

在声明计算之后,我们需要仔细地设计调度来更好地发挥性能。

调节 block/线程数的参数

# 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)

我们融合了 batch 矩阵相乘的外部维度,例如 op 维的 BB 和 FF 在 batch 矩阵相乘计算中通常称为「batch」维,我们用一个因子 (number_thread * vthread) 分割外部和内部维度。

在 batch 矩阵相乘中不需要 Strided 模式,因此将虚拟线程数(vthready 和 vthreadx)都设置为 1。

找到 number_thread 的最佳组合

下面的结果是基于 NVIDIA M40 GPU(CUDA 8.0)。

基于过去的经验,找到 num_thread_y 和 num_thread_x 最佳组合的方法是通过暴力搜索(brute-force search)。经过暴力搜索后,可以找到当前形状的最佳组合,在当前的计算中,num_thread_y = 8,num_thread_x = 32。

将 batch 矩阵相乘与其他运算融合

现有的「黑盒」cuBLAS 库调用一般会作为常用的「op 融合」优化策略的边界。然而,利用生成的高效 batch 矩阵相乘内核,融合边界极易被打破,将不仅仅是各个元素之间的融合,因此可以获得更好的性能改进。

从计算图中可以看出,batch 矩阵相乘之后总是伴随着广播加法运算或转置运算。

通过将「加法」或「转置」运算与 batch 矩阵相乘融合,可以减少内核启动开销和冗余内存访问时间。

batch 矩阵相乘和广播加法融合计算的声明如下:

# 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)

batch 矩阵相乘和转置融合计算的声明如下:

# 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 (separate): 541.9 us
  • TVM BatchMatmul: 37.62 us
  • TVM BatchMatmul + Transpose (fused): 38.39 us

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

集成 TensorFlow

在我们的工作负载中,batch 矩阵相乘的输入形状是有限的,易于提前枚举。有了这些预定义的形状,我们可以提前生成高度优化的 CUDA 内核(固定形状的计算可以带来最佳优化潜能)。同时,还将生成一个适合大多数形状的通用 batch 矩阵相乘内核,为没有提前生成内核的形状提供回退机制。

我们将生成的针对特定形状的高效内核和回退机制集成到 Tensorflow 中。我们开发了一些融合操作,例如 BatchMatMulTranspose 或 BatchMatMulAdd——使用 TVM runtime API 为确定输入形状启动特定生成的内核或调用回退内核。

通过执行图优化 pass,可以利用融合操作自动替换原始batch matmul + add/transpose。同时,通过结合更激进的图优化 pass,我们尝试利用 TVM 为长尾操作模式生成更高效的融合内核,以进一步提升端到端性能。

总结

在阿里巴巴,我们发现 TVM 是非常有效的开发高性能 GPU 内核的工具,可以满足我们的内部需求。

在本博客中,我们以 Transformer 模型为例,说明了我们利用 TVM 的优化策略。

首先,我们通过第一原理分析确定了 Transformer 模型的关键问题。然后,我们使用 TVM 生成高度优化的 CUDA 内核来取代 cuBLAS 版本(此时达到 13 倍的加速)。

接下来,利用 TVM 的内核融合机制来融合 batch 矩阵相乘的前/后操作,以带来进一步的性能改进(性能提升 1.7 倍)。端到端的性能改善达到 1.4 倍。基于这些生成的内核,我们开发了一个图优化 pass,自动用 TVM 融合内核替换掉原有的计算模式,确保优化对终端用户透明。

最后,所有这些优化都以松散耦合的方式集成到 TensorFlow 中,这展示了将 TVM 与不同深度学习框架集成的潜在方式。

目前我们还有一项正在进行的工作——将 TVM 整合为 TensorFlow 的 codegen 后端。我们希望将来与社群分享更多成果。

via:http://www.tvmlang.org

原文发布于微信公众号 - AI研习社(okweiwu)

原文发表时间:2018-03-27

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏深度学习计算机视觉

数据挖掘之数据预处理学习笔记数据预处理目的主要任务

数据预处理目的 保证数据的质量,包括确保数据的准确性、完整性和一致性 主要任务 数据清理 填写缺失的值、光滑噪声数据、识别或者删除离群的点,先解决这些脏数据,否...

2883
来自专栏CDA数据分析师

一篇文章教你如何用R进行数据挖掘

引言 R是一种广泛用于数据分析和统计计算的强大语言,于上世纪90年代开始发展起来。得益于全世界众多 爱好者的无尽努力,大家继而开发出了一种基于R但优于R基本文本...

3515
来自专栏目标检测和深度学习

教程 | 如何使用DeepFake实现视频换脸

1.3K2
来自专栏祁旭翔的专栏

【 SPA 大赛】win10 python3.5.X 下开启 lightgbm 支持

GBDT (Gradient Boosting Decision Tree) 是机器学习中一个长盛不衰的模型,其主要思想是利用弱分类器(决策树)迭代训练以得到最...

3.3K0
来自专栏鸿的学习笔记

写给开发者的机器学习指南(七)

Classifying email as spam or ham (NaiveBayes)

1211
来自专栏贾志刚-OpenCV学堂

基于积分图的二值图像膨胀算法实现

积分图来源与发展 积分图是Crow在1984年首次提出,是为了在多尺度透视投影中提高渲染速度。随后这种技术被应用到基于NCC的快速匹配、对象检测和SURF变换中...

5398
来自专栏人工智能LeadAI

如何使用sklearn进行数据挖掘

1.1、数据挖掘的步骤 数据挖掘通常包括数据采集,数据分析,特征工程,训练模型,模型评估等步骤。使用sklearn工具可以方便地进行特征工程和模型训练工作,在《...

2989
来自专栏量子位

亚马逊发布新版MXNet:支持英伟达Volta和稀疏张量

安妮 编译自 AWS官博 量子位 出品 | 公众号 QbitAI Apache MXNet v0.12来了。 今天凌晨,亚马逊宣布了MXNet新版本,在这个版本...

3956
来自专栏专知

【干货】RNN-LSTM的Keras实现:以预测比特币和以太坊价格为例(附代码)

【导读】本文是Siavash Fahimi撰写的一篇很棒的技术博文,主要讲解了用Keras实现RNN-LSTM,并用来预测比特币和以太坊的价格。在过去的一年,互...

6.9K9
来自专栏编程微刊

人工智能面试题86问,新手找工作必备!

3214

扫码关注云+社区

领取腾讯云代金券