2024 年 5 月，关于推荐模型自动编译优化的论文《RECom: A Compiler Approach to Accelerate Recommendation Model Inference with Massive Embedding Columns》在系统领域顶会 ASPLOS 2024 上中选并进行了展示，并被授予了 Distinguished Artifact Award 荣誉，以表彰 RECom 的易用性与结果的可复现性。

RECom 是一个旨在加速 GPU 上深度推荐模型推理的机器学习编译器，解决了已有机器学习编译器无法有效优化推荐模型中重要的嵌入层计算的问题。RECom 提出了一种新颖的面向子图间并行的融合方法，以在单个 GPU 内核中生成处理大量嵌入列的高效代码，大幅减少了内核启动的开销并显著提升了 GPU 的利用率。通过一个形状计算简化模块来解决动态形状下形状计算带来挑战，并通过一个嵌入列子图优化模块来消除冗余计算。在阿里巴巴的四个真实的内部生产推荐模型和两个生成模型上评估了 RECom，实验表明，RECom 在端到端推理延迟和吞吐量方面分别优于 TensorFlow 基线 6.61 倍和 1.91 倍。

背景

基于深度学习的推荐模型在各大企业的业务中变得越来越重要。通过为用户推荐合适的商品，能够帮助提升客户体验，增加销售额，以及提高客户保留率。典型的推荐模型包括两部分，即嵌入层和深度神经网络层（DNN）。 在生产环境中，嵌入层通常由对应不同特征字段的大量嵌入列（即通过查表将输入特征转换为嵌入向量的子图，如下图中虚线方框所示）组成，它能将输入的特征（比如用户 ID，商品 ID 等）映射到低维的连续向量。

为了获得更高的模型精度，开发人员通常会生成数千个统计特征，并使用不同的嵌入列来处理它们。 然而，处理如此大量的嵌入列是昂贵的。 例如，我们在阿里巴巴的模型上进行的实验表明，嵌入列可以占据 GPU 上 99%以上的端到端推理延迟。 因此，业务中迫切需要对推荐模型中的海量嵌入列进行优化。

挑战

虽然目前已有一些加速嵌入层计算的手写算子库[1]，但它们很难在真实业务中被广泛应用，因为它们不可能在库中遍历嵌入层中所有可能的算子组合。此外，出于隐私考虑，实际业务中经常要求根据模型的计算图 IR（比如 TensorFlow GraphDef）而非源代码来进行优化。然而，一个公司的业务中通常会包含大量具有不同嵌入列结构的模型，并且一个模型可以包含上百万行的 IR。因此，基于这些 IR 再手工地使用这些算子库来重新搭建所有的业务模型是不现实的。

另一方面，机器学习编译器，如 XLA[2]等被广泛地应用于自动加速机器学习模型的计算。然而，目前的机器学习编译器主要关注于 DNN 的计算，而无法为嵌入层中包含的大量嵌入列提供有效的优化。它们主要存在以下三个问题：

首先，在 GPU 上为数千个嵌入列中的大量算子生成有效的代码是一项挑战。现有编译器的算子融合策略既无法有效消除模型中大量的非计算开销（包括内核启动与算子调度），也无法充分利用模型中子图间的并行性。我们在一个特定模型的实验上发现，XLA 可以为该模型生成超过 10,000 个 GPU 内核，这会引入大量的非计算开销，进而导致 GPU 的活跃时间仅占 33%。并且，线上推理时输入的 batch size 通常不会太大，导致大多数生成的 GPU 内核的 Waves per SM 都非常小（小于 0.06），意味着大部分的 GPU SM 在执行中都没有被利用。

其次，推荐模型通常具有动态形状（dynamic shapes）。传统的机器学习编译器通常依赖张量的形状来进行各种优化，而推荐模型的计算图中的张量形状一般无法在编译期获取。一方面，动态形状造成的张量间形状关系信息缺失阻碍一些编译优化的 pass。另一方面，动态形状场景下，部分算子（如 SparseReshape）需要形状的输入，继而引入了大量形状计算。形状计算与张量计算的耦合使得计算图变得复杂，阻碍了后续的图优化。

下图(a)展示了一个级联 SparseReshape 的子图，其中蓝色的算子表示形状计算。由于 SR1 输出的 sparse tensor 的消费者只有 SR2，因此第一次 reshape 是冗余的。理想情况下，编译器检测到这样的情况后可以直接把 SR1 的输入张量（即张量 1）连接到 SR2，越过冗余的 SR1，如(b)中所示。然而，由于形状计算与张量计算的耦合，导致实际转化后的图会如(c)中所示，SR1 并不能被正确的优化掉。

第三，现有的机器学习编译器忽视了推荐模型中存在的大量冗余计算。通过对业务模型的分析，我们观察到现有模型普遍存在大量的冗余计算。例如，使用 TensorFlow 框架进行嵌入表查找时会引入一些边界检查的计算，而这些检查在整个嵌入列子图的上下文里可能是冗余的，因为在经过哈希/分桶等操作后，用于嵌入表查找的索引已经满足了要求。在极端情况下，一个嵌入列中 80%的 GPU 计算时间可能都是冗余的。

破局

我们提出了第一个端到端的推荐模型优化编译器 RECom 来解决以上的问题。

针对大量嵌入列的 GPU 代码生成，RECom 提出了一种子图间并行性导向的算子融合方法，将嵌入层中数千个嵌入列包含的大量算子全部融合（fuse）到一个 GPU kernel 中。一方面，这种方法使得 RECom 能够消除频繁的内核启动与框架调度，显著减少嵌入层执行时的非计算开销。另一方面，为了生成具有高并行性的代码，RECom 将独立的嵌入列映射到不同组的 GPU 线程中。每个线程组中的线程协作处理一个嵌入列，并利用 GPU 的层次存储来进行中间数据缓冲和线程间通信。通过这种方式，RECom 可以有效地利用 GPU 来挖掘推荐模型嵌入层中算子内部和嵌入列间的并行性。

为了解决动态形状与形状计算带来的问题，RECom 提出了一种基于符号表达式的方法。类似于 BladeDISC[3]，RECom 通过下图中的四种 pattern 构建了嵌入列的全局符号形状表达式，并利用这些表达式来确定不同张量的形状间关系。

得到形状的符号表达式后，RECom 进一步将所有形状计算子图用统一的算子进行重构。重构后的 ShapeConstruct 算子仅依赖于生成所需符号的算子，而非依赖于原计算子图的输入。这样，RECom 得以解耦张量计算与形状计算来简化计算图。回到之前冗余 SparseReshap 的例子，再经过形状计算重构后，SR2 输入形状对 SR1 的依赖被消除，因此在经过原来的编译优化 pass 后，我们能得到下图的结果，成功去除了冗余的 SR1。

对于嵌入层中存在的大量冗余计算，我们在 RECom 中开发了一个嵌入列子图优化模块，针对性地为常见的冗余计算设计了优化。具体细节可以参考我们的论文。

RECom 的整个编译优化流程如下图所示。

与此同时，阿里云 PAI 团队与悉尼大学合作的论文《MonoNN: Enabling a New Monolithic Optimization Space for Neural Network Inference Tasks on Modern GPU-Centric Architectures》在 OSDI2024 中选，论文创新性地提出了 MonoNN，这是第一个能够为单个 GPU 上的常见静态神经网络（NN）推理任务提供新的整体设计和优化空间的机器学习优化编译器。MonoNN 可以将整个神经网络容纳到单个 GPU 核心中，大大减少了非计算开销，并从新形成的整体优化空间提供进一步细粒度的优化机会。最重要的是，MonoNN 确定了各种 NN 运算符之间的资源不兼容性问题，这是创建这样一个整体优化空间的关键设计瓶颈。MonoNN 通过系统地探索和利用跨不同类型的 NN 计算的并行补偿策略和资源权衡，并提出了一种新的独立于调度的群体调谐技术来显著缩小极其庞大的调谐空间。评估表明，MonoNN 的平均加速比比最先进的框架和编译器提高了 2.01 倍。具体而言，就端到端推理性能而言，MonoNN 在 TVM、TensorRT、XLA 和 AStitch 方面的表现分别高达 7.3 倍、5.9 倍、1.7 倍和 2.9 倍。目前，MonoNN 已开源，如需获取 MonoNN 源代码请前往https://github.com/AlibabaResearch/mononn;

论文地址：https://dl.acm.org/doi/10.1145/3623278.3624761

论文信息

论文标题：《RECom: A Compiler Approach to Accelerating Recommendation Model Inference with Massive Embedding Columns》

论文地址：https://dl.acm.org/doi/10.1145/3623278.3624761

参考文献

[1] NVIDIA. NVIDIA-Merlin/HugeCTR. https://github.com/NVIDIA-Merlin/HugeCTR, 2023.

[2] Google. XLA: Optimizing Compiler for Machine Learning. https://www.tensorflow.org/xla, 2023.

[3] Zhen Zheng, Zaifeng Pan, Dalin Wang, Kai Zhu, Wenyi Zhao, Tianyou Guo, Xiafei Qiu, Minmin Sun, Junjie Bai, Feng Zhang, Xiaoyong Du, Jidong Zhai, and Wei Lin. BladeDISC: Optimizing Dynamic Shape Machine Learning Workloads via Compiler Approach. Proc. ACM Manag. Data, 1(3), nov 2023.