数据工程与知识工程教育部重点实验室张峰教授与阿里巴巴集团的合作研究成果，在2024年度的ACM架构支持编程语言和操作系统国际会议（ASPLOS）上被授予了Distinguished Artifact Award，以表彰其研究成果的易用性与可复现性。该研究成果的论文题为《RECom：利用编译器加速具有大规模嵌入列的推荐模型推理》。该论文以阿里真实应用为基础提出了一种新的机器学习编译器方法，优化了GPU上复杂推荐系统的计算过程。

ASPLOS是CCF A类计算机系统研究的顶级国际会议，涵盖计算机体系结构、编程语言、操作系统及其相关领域如网络和存储的多学科内容，与ISCA、MICRO、HPCA并称为体系结构四大顶会。此次会议在2024年4月27日至5月1日在美国圣地亚哥举行。

研究摘要

基于深度学习的推荐模型在各大企业的业务中变得越来越重要。通过为用户推荐合适的商品，它们能够帮助提升客户体验，增加销售额，以及提高客户保留率。嵌入层对于保证深度推荐模型的准确性至关重要，它能够将输入特征转换为数值的嵌入向量。然而，由于现代推荐模型经常包含数千个嵌入列来处理相应的特征字段，嵌入层的执行过程通常非常耗时。另一方面，包括阿里巴巴在内的许多公司在业务中广泛使用机器学习编译器来自动优化机器学习模型。然而，现有的编译器主要关注传统深度神经网络结构的优化，忽略了深度推荐模型为编译优化带来的一系列挑战。本文提出了RECom，这是第一个旨在优化推荐模型中嵌入层的机器学习编译器。RECom解决了三个主要挑战。首先，在GPU上为大量嵌入列内的大规模运算符生成高效的程序是困难的。其次，动态形状场景下的复杂形状计算阻碍了进一步的图优化。最后，机器学习框架不可避免地引入了由于稳健性考虑而产生的冗余计算。实验表明，RECom在终端推理延迟和吞吐量方面，相比于现有最先进的基线方法，分别实现了6.61倍和1.91倍的性能提升。系统已在AlibabaResearch开源，并通过了ASPLOS全部Artifact Evaluation，荣获Distinguished Artifact Award。欢迎使用：

https://github.com/AlibabaResearch/recom。

详细内容

近年来，深度推荐模型已在不同的商业场景中得到广泛应用。其中，Google YouTube利用深度推荐模型进行视频排名，Meta也在在线广告应用中采用了类似的方法，而阿里巴巴则将其应用于电子商务领域。根据 Meta 的报告，深度推荐模型已成为机器学习工作负载中不可或缺的一部分。据报告显示，它们在数据中心总体机器学习训练和推理的周期中分别贡献了超过 50%和 70% 的比例。

典型的深度推荐模型包括两部分，即嵌入层和深度神经网络层（DNN）。 在生产模型中，嵌入层通常由对应不同特征字段的大量嵌入列（即通过查表将输入特征转换为嵌入向量的子图）组成，它能将输入的特征（比如用户ID，商品ID等）映射到低维的连续向量。为了获得更高的模型精度，开发人员通常会生成数千个统计特征，并使用不同的嵌入列来处理它们。 然而，处理如此大量的嵌入列是昂贵的。 例如，我们在阿里巴巴的模型上进行的实验表明，嵌入列可以占据GPU上99%以上的端到端推理延迟。因此，业务中迫切需要对推荐模型中的海量嵌入列进行优化。

与此同时，机器学习编译器被广泛用于自动优化传统的机器学习模型。通常来说，机器学习框架（比如TensorFlow）将一个机器学习模型的执行过程描述为一个有向无环的数据流图，即计算图。计算图中的每一条边都对应一个张量的流动，而每一个节点则表示对输入的张量应用一个特定的算子，如矩阵乘法、ReLU等。机器学习编译器可以通过一系列的图变换来优化模型的计算过程，并根据计算图在给定的硬件（比如GPU）上自适应地生成高性能的程序。然而，目前的机器学习编译器都无法有效地优化嵌入列。在本文中，我们重点关注使用编译器通过GPU上的大量嵌入列来加速推荐模型推理。

开发机器学习编译器来加速GPU上深度推荐模型的推理非常重要，因为它可以帮助公司节省大量的手动工作。传统上，许多公司聘请专家团队来手工优化其推荐模型。然而，一个公司的不同业务线可能有数千个深度推荐模型，具有不同的嵌入列结构。此外，出于隐私考虑，许多实际业务需要基于计算图的中间表示（IR）而不是源代码来优化模型。例如，阿里巴巴的一个工业模型包含350万行TensorFlow GraphDef格式的IR。通过分析如此大量的IR来手动重写模型是不切实际的，因此需要机器学习编译器方法来自动优化模型。

推荐模型中嵌入列的特性给编译优化带来了三个基本挑战。首先，在GPU上为数千个嵌入列中的大量算子生成有效的调度是一项挑战。传统的解决方案无法消除大部分大量的非计算开销，也无法充分利用子图间的并行性。对于特定模型，XLA可以生成超过10,000个内核，这会引入大量的非计算开销，然后导致GPU活动时间仅占33%。此外，大多数生成的内核的Waves per SM都非常小（小于0.06）。其次，推荐模型的动态形状特征引入了许多形状计算操作。形状和张量计算的混合使得图拓扑过于复杂而无法优化。第三，出于鲁棒性考虑，框架可以引入许多冗余操作。我们发现，在极端情况下，冗余计算可以占GPU上整个嵌入处理时间的80%。

在我们的合作AIR项目中，我们联合提出了RECom，这是一个旨在加速GPU上深度推荐模型推理的机器学习编译器，它有效地解决了上述挑战。首先，我们开发了一种新颖的面向子图间并行的融合方法，以在单个内核中生成处理大量嵌入列的高效代码。它识别与每个嵌入列对应的子图，将每个子图映射到GPU上的一组线程，然后利用GPU层次化的存储结构将子图内的计算融合在一起。其次，我们开发了形状计算简化模块来解决形状计算挑战。它将所有张量形状表示为符号表达式，并根据符号表达式重建所有复杂的形状计算。第三，为了消除冗余计算，我们开发了一个嵌入列子图优化模块，通过分析计算图的上下文来执行简化。

我们在阿里巴巴的四个真实的内部生产推荐模型和两个生成模型上评估了RECom。实验表明，RECom在端到端推理延迟和吞吐量方面分别优于最先进的基线方法6.61倍和1.91倍。总之，这项工作做出了以下贡献：

l 我们揭示了推荐模型推理期间嵌入计算的三个重大性能挑战，并提供了一组解决方案和见解。

l 我们联合提出了RECom，这是第一个优化推荐模型推理过程中耗时的嵌入列计算的机器学习编译器。

l 我们在四个工业推荐模型和两个综合推荐模型上评估RECom。实验结果表明 RECom 显着优于基线。

潘再峰：中国人民大学计算机系硕士生，研究兴趣为机器学习系统与编译器，将于2024年秋天前往美国加州大学圣地亚哥分校攻读博士学位。