阿里巴巴NMT团队：TVM优化Transformer在GPU推理中的速度提升

阿里巴巴机器翻译团队在推动技术前沿，针对其神经机器翻译(NMT)服务的性能优化工作。NMT是一种端到端的自动翻译方法，相比传统的短语翻译系统，具有显著的优势。阿里巴巴已将Transformer[1]作为NMT系统的基石，因其在离线训练阶段表现出高准确性和效率。然而，在在线推理环节，Transformer的性能受限，尤其是在GPU上，初始版本的推理速度相比LSTM模型稍显逊色，尤其是在1.5X到2X之间。 为了解决这一问题，团队采用了一种创新的方法，即引入TVM（Tensorflow虚拟机）来优化Transformer模型的GPU运算。TVM是一个高性能的深度学习计算框架，它能够生成针对特定硬件的优化内核。在Transformer模型中，批量矩阵乘法（batch matmul）是关键性能瓶颈，特别是在多头注意力计算中，它占据了GPU执行时间的约30%。 原始cuBLAS库中的批量matmul实现并未达到最佳优化，这使得团队决定使用TVM进行定制化优化。通过TVM，他们实现了至少13X的加速，显著提升了批量matmul的计算效率。这不仅提高了Transformer的整体推理速度，还为其与其他系统和服务的集成提供了更快的响应。 批量matmul在Transformer中的应用涉及矩阵与矩阵的乘法，其中所有实例共享相同的维度规格，这对于并行计算非常有利。在语言翻译任务中，其形状通常较小，如输入句子长度和解码器步长限制，一般小于30。TVM的介入使得团队能够针对性地优化这些计算密集型操作，从而极大地提高了NMT在实际生产环境中的性能表现。 通过引入TVM并优化关键计算环节，阿里巴巴机器翻译团队成功地优化了GPU上神经机器翻译的效率，为全球电子商务的自动化翻译服务奠定了坚实的技术基础，也为其他深度学习应用的硬件优化提供了一个有价值的实践案例。