阿里巴巴最新实践:TVM+TensorFlow优化GPU上的神经机器翻译

本文是阿里巴巴 PAI-Blade 团队发表于 TVM 的最新博文,文中阐述了如何将 TVM 引入 TensorFlow,使 TensorFlow 中的 batchmul 速度提高 13 倍,同时将端到端神经机器翻译性能提高 1.7 倍。

背 景

神经机器翻译(NMT)是一种端到端的自动翻译方法,有可能克服传统短语翻译系统的不足。最近,阿里巴巴集团正在尝试将 NMT 服务部署到全球电子商务业务场景中。

目前,我们正在尝试开发 Transformer[1] 作为我们 NMT 系统的核心组件,因为它比基于 RNN/LSTM 的经典模型更有利于高效的离线训练,具有同等(甚至更高)的精确度。虽然 Transformer 打破了时间步长之间的依赖关系,因此对于离线训练阶段非常友好,但是对于在线推理来说,就没有那么高效了。在我们的生产环境中,初始版本的 Transformer 推理速度比 LSTM 版本慢大约 1.5 倍到 2 倍。为了提高推理性能,我们已经尝试进行了多种优化,如图级别 op 融合、循环不变式节点运动 [3]。我们发现批量 matmul 是 Transformer 中的一个性能关键点,而当前 cuBLAS 的实现并未进行很好的优化。

我们的实践结果表明,通过 TVM 生成的内核(优化了 schedule )可以使批量 matmul 计算速度提高至少 13 倍,并且在启用了算子融合的情况下还将进一步提高速度。

批量 Matmul

为什么要使用批量 matmul

在 Transformer 中,批量 matmul 被广泛应用于多头 attention(multi-head attention)的计算。利用批量 matmul,attention 层中的多头可以并行运行,有助于提高硬件的计算效率。

我们在推理阶段对 Transformer 模型进行了深入的剖析,结果表明批量 matmul 计算占 GPU 内核执行时间的 30% 左右。我们利用 nvprof[2] 对 cuBLAS 的批量 matmul 内核进行了第一性原理分析,结果表明当前的实现效果很差,同时我们还发现了一些有趣的现象。

什么是批量 matmul

通常,批量 matmul 计算是指在一批矩阵上执行矩阵 - 矩阵乘法。其中 batch 被认为是“均匀的”,即所有实例对于它们各自的 A、B 和 C 矩阵具有相同的维度(M、N、K)、前导维度(lda、ldb、ldc)和转换率。

批量 matmul 计算可以更具体地描述如下:

批量 matmul 的 shape

在语言翻译任务中,批量 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, 最大编码长度] 的范围内。

cuBLAS 批量 matmul 的性能问题

首先,我们对批量 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 内核。

批量 matmul 计算

在 TVM 中,通常批量 matmul 计算可以声明为:

Schedule优化

在对批量 matmul 计算进行声明之后,我们需要仔细设计Schedule,以充分挖掘其性能潜力。

调整块/线程号的参数

我们融合了批量 matmul 的外部维度,即 op 维度的 BB 和 FF,在批量 matmul 计算中通常称为“批量”维度。然后我们用因子(number _ thread * vthread)分割外部维度和内部维度。

批量 matmul 中不需要 stridged 模式,因此虚拟线程数(vthread_y 和 vthread_x)都设置为 1。

寻找 number_thread 的最佳组合

下图是使用 CUDA8.0 在 NVIDIA M40 GPU 上得到的结果。

根据以往的经验,寻找 num_thread_y 和 num_thread_x 最佳组合的方法是暴力搜索。通过暴力搜索,可以找到当前形状(shape)的最佳组合,我们在当前计算中得到的最佳组合为 num_thread_y=8,num_thread_x=32。

将批量 matmul 与其他运算融合

当前的“黑盒”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 倍。

与 TensorFlow 集成

在我们的工作负载中,批量 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


原文发布于微信公众号 - 奇点(qddata)

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

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

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏ATYUN订阅号

使用Python进行人脸聚类的详细教程

思考下面这个场景:两名劫匪在抢劫波士顿或纽约等繁华城市的银行。银行的安全摄像头工作正常,捕捉到了抢劫行为,但劫匪戴着头套,没办法看到他们的脸。

5882
来自专栏PPV课数据科学社区

大数据竞赛平台-Kaggle入门篇

本文作者: wopon_ 来源:36大数据 本文长度为1500字,建议阅读4分钟 这篇文章适合那些刚接触Kaggle、想尽快熟悉Kaggle并且独立完成一个竞赛...

4778
来自专栏人工智能

使用10几行Python代码,快速建立视觉模型识别图像

视觉 进化的作用,让人类对图像的处理非常高效。 这里,我给你展示一张照片。 ? 如果我这样问你: 你能否分辨出图片中哪个是猫,哪个是狗? 你可能立即会觉得自己遭...

4009
来自专栏人工智能头条

详解TensorBoard如何调参

1023
来自专栏ATYUN订阅号

自相关与偏自相关的简单介绍

自相关和偏自相关图在时间序列分析和预测中经常使用。这些图生动的总结了一个时间序列的观察值与他之前的时间步的观察值之间的关系强度。初学者要理解时间序列预测中自相关...

6574
来自专栏大数据挖掘DT机器学习

Python机器学习库scikit-learn实践

一、概述 以最广泛的分类算法为例,大致可以分为线性和非线性两大派别。线性算法有著名的逻辑回归、朴素贝叶斯、最大熵等,非线性算法有随机森林、决策树、神经网络、核...

3805
来自专栏人工智能

完全云端运行:使用谷歌CoLaboratory训练神经网络

选自Medium 作者:Sagar Howal 机器之心编译 参与:路雪 Colaboratory 是一个 Google 研究项目,旨在帮助传播机器学习培训和研...

8008
来自专栏新智元

【重磅】Jeff Dean等提出自动化分层模型,优化CPU、GPU等异构环境,性能提升超 60%

谷歌大脑Jeff Dean等人最新提出一种分层模型,这是一种灵活的端到端方法,用于优化CPU、GPU等的自动化设备配置。该方法在多个主要神经网络模型上测...

3367
来自专栏从流域到海域

A Gentle Introduction to Autocorrelation and Partial Autocorrelation (译文)

A Gentle Introduction to Autocorrelation and Partial Autocorrelation 自相关和偏自相关的简单...

3036
来自专栏人工智能

TensorFlow实战——图像分类神经网络模型

Learn how to classify images with TensorFlow 使用TensorFlow创建一个简单而强大的图像分类神经网络模型 by...

4136

扫码关注云+社区

领取腾讯云代金券