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

最新推荐文章于 2024-06-25 15:54:43 发布

wujianming_110117

最新推荐文章于 2024-06-25 15:54:43 发布

阅读量846

点赞数

分类专栏： GPU，NPU，XPU，DPU，MPU 视频推理智能推理

吴建明

本文链接：https://blog.csdn.net/wujianing_110117/article/details/114514035

版权

本文探讨了如何使用TensorFlow和TVM来优化神经机器翻译（NMT）中的Transformer模型，特别是针对批量矩阵乘法（batch matmul）的性能瓶颈。通过TVM生成和调度优化内核，实现了13倍的加速，并通过融合运算符进一步提升了1.7倍性能。这些优化已集成到TensorFlow中，以透明地提高端到端推理效率。

摘要由CSDN通过智能技术生成

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