小红书训推异构引擎的设计与应用

原创曾鸣堃 DataFunSummit

2024年08月12日 10:02

导读今天给大家带来的分享主题是小红书训推异构引擎的设计与应用。

主要内容包括以下五大部分：

1. 小红书模型工程面临的挑战

2. 异构弹性引擎的设计与实施

3. 面向未来 HPC 训练框架

4. AI 编译技术

5. 展望未来

分享嘉宾｜曾鸣堃小红书训推计算引擎负责人

编辑整理｜娄政宇

内容校对｜李瑶

出品社区｜DataFun

小红书模型工程面临的挑战

近年来小红书业务一直保持高速增长，日均曝光、互动 UV 和日均搜索量都达到了非常高的水平。

从模型工程的角度来看，面临的主要挑战包括：

首先，随着模型的复杂度增加，其处理的数据量也越来越大；
其次，计算流程的需求也会随之增长；
此外，模型的应用场景也越来越广泛，不仅限于传统的搜索广告和推荐，还包括电商和直播等新业务。

面对业务高速增长带来的挑战，我们需要通过引擎技术有效降低成本，并为业务提供迭代空间。

02 异构弹性引擎的设计与实践

1. 第一代训练框架

接下来介绍一下训练框架的演进过程。

第一代训练框架 Larc 是一个基于 PS worker 的设计，主要包含以下几个核心技术点：

第一个是它能够支持超大规模的稀疏特征。
第二个是区别于传统 TensorFlow 的 Feature Column 的 Embedding Table 实现，我们使用了基于无冲突哈希表的实现，可以确保每个 ID 都能得到充分的学习。
第三是有一个高性能的 Lookup Table 算子实现。从上图中可以看出，我们将整个稀疏 Embedding 操作抽象成了几个算子。与原始的 TensorFlow 相比，我们进行了许多 OP 的融合，从而获得了更高效的性能收益。
第四个是支持了多种的参数的优化器。

如上图是小红书第一代 GPU 训练框架，可以看到，这里的核心设计点就是把 worker 作为一个 GPU 的 node。小红书的模型场景非常多，导致了我们的模型类型会非常多。另外我们大量使用了云服务，云厂商提供的 GPU 机型也会非常多。第一代框架的核心技术点是把一些计算密集的算子 place 到 GPU 执行，同时也做了一些 Lookup 的异步化的优化。

因为我们有大量稀疏特征，这导致 CPU 占用率很高，但是 GPU 利用率很低。所以这一代框架的核心问题是无法高效利用异构算力的优势。为了解决这个问题，我们设计了新一代的基于 CPU、GPU 的异构的训练框架。

2. 异构 GPU 训练框架

这里把整个搜广推的模型训练抽象成了十个步骤，我们调研了业界的一些方案，主要有两个方案：一个是在特征抽取之后去做拆分，一个是在 embedding 的序列聚合之后做拆分。结合我们的实际场景来看，方案二是更适合的，所以我们最后采用了方案二，把整个计算图的十个步骤一拆为三。

接下来介绍一些系统实现的细节。我们在异构训练框架上，把稀疏的网络子图，包括样本解析、特征抽取和 embedding Lookup 以及 embedding 的聚合这些操作放在 CPU worker。像稠密的网络子图、Dense 网络计算这部分，转移到 GPU worker 去算。

在设计细节上，我们采用了与 TensorFlow 相似的设计哲学。我们使用了成对的 Send/Receive OP，使得 Tensor 之间可以得到充分的并发通信。但与 TensorFlow 不同的是，我们通过插入大量的消息队列，使得 CPU worker 和 GPU worker 之间可以得到异步的解耦。这使得 CPU 和 GPU worker 可以充分地异步计算，充分利用硬件的计算能力。这也引出了第二个问题，即如何高效地匹配 N 个 CPU worker 和 M 个 GPU worker。为了解决这个问题，我们设计了一套基于全局动态分发的队列，其中只存储一些较小的 Meta 信息。通过这种方式，CPU 和 GPU worker 之间得到了高效的匹配。

第三点，我们在 CPU 和 GPU 的通信中间做了大量的混合精度的工作。使得 GPU 的入口带宽得到大幅度降低。从设计的架构图中可以看出，我们的训练角色被拆分为 CPU worker、CPU worker 和 PS 三类角色。对于一些高性能和高核卡比的机型，它们在拆分后可能会出现 GPU 上的 CPU 相对空闲的情况。针对这类 case，我们会部署一个 local 的 CPU worker。通过这种方式，可以使得 GPU 硬件上的 CPU 和 GPU 都得到高效使用。目前这套框架已经覆盖公司内核心场景的模型训练，使训练速度得到了显著提升。

3. 第一代 GPU 推理架构

接下来介绍推理架构的演进。第一代 GPU 推理架构如上图所示，可以抽象成三个主要步骤，即特征抽取、TF serving 和多目标融合的 ValueModel 的计算。

这一代 GPU 推理架构的主要问题和训练相似，存在 GPU 硬件利用率较低的问题，其瓶颈在于 CPU 利用率较高。另外，迭代不够灵活，算法迭代时只能去加 GPU 机器，即便此时 GPU 并没有被打满。还有一个问题是小 batch size，这是推理场景都会遇到的问题，因为推理场景的 batch size 一般会比训练场景要低一个数量级。小 batch size 对 GPU 的利用率较低。

4. 异构 GPU 推理架构

为解决这些问题，我们设计了新一代的异构拆分的推理架构。核心工作是将整个 TensorFlow 的图拆分为 CPU worker 和 GPU worker 两个部分。这里面临的挑战主要包括：首先，在拆分更多在线模块后，在线延迟必然会上涨；其次，小 batch size 无法充分利用 GPU 的计算单元；最后，如何精细拆分有状态的服务，如初排服务，以实现计算和内存的解耦。

我们需要解决的第一个挑战就是模块拆分引入 latency 上涨的问题。我们发现，在我们的系统中，latency 上涨主要来自于一些 IO 密集型的数据处理任务。在我们的场景中，每个请求携带的数据量差异很大。因此，为了有效地控制 latency，我们设计了一个基于数据包大小动态并行的策略，这可以有效降低大数据包带来的长尾 latency 问题。

第二个优化是针对模块拆分引入的序列化和反序列化这两个阶段，做了 zero copy 的优化策略。在优化前，有数据的行列转换，以及数据的序列化这样两次的低效操作，现在这两部分都被优化掉了。

为了解决小 batch size 情况下对 GPU 计算单元利用不充分的问题，我们设计了 auto batching 的策略。设计思路是在同一时间窗口内将多个请求聚合成一个请求进行计算。从请求的角度来看，单个请求的 batch size 会变大，从而充分利用 GPU 的计算单元。从 kernel 的角度来看，通过请求聚合，可以同比例地降低 kernel 的执行次数和 overhead 的开销。整个 kernel 的执行分为 launch 和 kernel 执行两步，通过请求聚合，会变成一个 kernel 去执行，从而降低对系统的 overhead。

关于如何优化有状态的服务，我们调整了线上初排服务的设计。在左边的图中可以看到，初排服务在整个检索链路中是一个承上启下的模块，因此需要做大量的缓存设计，通过牺牲一些数据和模型的时效性来预估更多的笔记数量。在这张图中，我们也可以看到初排模块有很多内存密集型和计算密集型的子模块。如果我们想要对初排模块进行异构拆分，让它们之间得到更好的解耦，这是一个很大的挑战。

右图是我们的解决方案，主要设计思想是让尽可能多的数据在一个设备内闭环，避免跨网络通信开销。因此，我们设计了一个三次请求的协议，并最终实现了计算和内存之间的充分解耦。

然而，这三次请求的发生阶段不同，因此对工程提出了挑战。因为有三次请求，第二次和第三次请求必须在同一台节点上执行。此外，三个请求可能会超时，因此，我们需要在 GPU 机器上进行状态的 TTL 管理。

新的异构推理引擎在几个核心主场景落地后，高峰期的 GPU 利用率都提高了 100% 以上。

5. 总结

异构训练推理引擎带来了以下几大优势：

首先是高性能，成本显著低于上一代引擎，CPU 和 GPU 利用率都能达到 65% 到 95%。

第二是灵活性好，算法可以灵活地根据机型和模型选择最优的算力拆分策略。

第三是迭代 ROI 高，业务可以根据区域和模型增加的消耗来申请资源。

第四是伸缩性好，因为训练和推理业务在业务场景上天生就比较适合做潮汐混部，而且我们目前使用的机型也是一致的，因此这套架构能够很好地支持这种需求。

03 面向未来的 HPC 训练框架

第三部分将介绍面向未来的 HPC 训练框架。

未来的一个趋势是业务高速增长且算法快速迭代，我们将面临更大规模的数据集，以此倒推，我们需要更高的训练吞吐量以确保训练效率。其次，模型的复杂度也在增加。近年来，一些常见的复杂网络建模方法，如 PPNet、PLE 等，以及最近流行的生成式推荐系统方向，都将使得 Dense 参数变得更加复杂。因此，我们将面临更大的计算需求。

现有训练框架存在的主要问题包括：

首先，无论是异构的训练引擎还是之前的 Larc，都是基于 PS worker 架构。然而，随着节点数量的增加，训练加速比会有显著的衰减。
其次，为了提高训练吞吐量，我们会增加更多训练节点。但由于 PS worker 采用异步训练模式，增加更多节点会导致更大的异步性，从而影响模型的整体收敛效果。
此外，像 A10、A30 这样的设备虽然基于 PCIE 和 TCP 网络通信，但难以高效地支持大规模复杂 Dense 模型的训练。

为了解决这些问题，我们开发了一种面向未来的新一代 HPC 训练框架。该框架借鉴了百度 AIBox 和 NVIDIA HugeCTR 等框架的实现原理。

虽然我们使用了基于 A100 和 H800 的大算力 GPU，但 A100 的显存只有 640GB，而我们的模型规模已经达到了好几 TB。因此，将所有模型全部放入显存是不现实的。此外，在通信方面，整个模拟训练过程中需要大量的通信和计算，如何高效地处理两者之间的关系是一个挑战。最后，由于我们的稀疏特征非常丰富，有几十到上百个，因此如何高效地完成在 GPU 内的 embedding 相关操作也是一大挑战。

上图展示了我们的 HPC 训练架构的设计。由于场景的稀疏性导致算力消耗较大，而 HPC 机器上的 CPU 核心数量也不多。因此，如果不进行异构拆分，CPU 的算力将无法得到充分利用。因此，我们将一些 IO 密集型的任务，如样本解析和特征处理，拆分到 CPU 集群中执行。右侧是我们的 HPC 集群，它会接收特征抽取后的特征并进行模型后续计算。这里分为三个角色：GPU worker 主要负责图内的稠密计算；有两层的 PS，包括 HBMPS 和 DRMPS。

与业界不同的一些核心技术点包括：

首先，我们采用了大规模的 pass 粒度聚合，利用样本间数据的局部性原理，进行了大批量的 pass 后进行统一的去重，从而减少了单位训练所需的 ID 数量。
其次，我们进行了 embedding 的置换，因为我们无法将所有模型都放入显存中，因此需要进行批量置换。在置换过程中，如果采用最原始的方式，即与图内模型训练串行进行，会导致整个 GPU 的利用率非常低。因此，我们进行了 embedding 的置换和图内计算的流水线并行优化，使得图内计算和 embedding 的置换之间可以解耦。
第三，我们做了增量式换入换出，因为原始实现中每轮都会全量换入和换出，但 ID 量很大，全量换入换出会花费很多时间，降低模型训练效率。因此，我们利用相邻两个 pass 之间的数据局部性来提高效率。每次只换入增量的部分。也就是说，在第二个 pass 换入时，会查看 GPU 内部是否已经存在该部分，如果不存在，则只换入这部分。其次，当我们在多个 pass 中积累了足够的数据后，会将后面几轮用不到的 ID 替换掉。这样可以提高整个 ID 的换入和换出效率。
最后是 Table Fusion，与之前提到的方法类似，由于小红书的原始 embedding 列非常多，因此我们按照 embedding 维度进行了一层聚合，从而减少了图内算子的数量，提高了整体模型训练的效率。

04 AI 编译技术

1. AI 编译技术

AI 编译技术要解决的问题是从任何水位的模型中，通过AI 编译技术得到一个相对下限比较高的模型性能表现。整个 AI 编译器技术栈可以分为前端、中端和后端三个部分。前端主要包括自研的 Fourier 图优化器和 TensorFlow Grappler 优化器。中后端主要包括业界常用的 XLA、TVM、TensorRT 等。

2. 前端优化

关于前端优化的工作，首先来介绍 Fourier 优化器，它是一种 rule-based 对子图高效匹配和改写的方法。此外，TensorFlow 中也有一个名为 Grappler 的前端图优化器。然而，它的迭代数较少，策略也比较通用且不够丰富，Fourier 对此进行了补充。针对搜广推场景的高频子图做了识别和高效的替换改写。因为搜广推场景中的粗排精排会大量用到多目标的建模，会面临数据 tensor 的切分和并行的执行计算，这些非常适合通过改写后做一些算子的融合，来提高整体的计算效率。

Fourier 主要做了两个方面的优化：

第一个是核心的低效子图识别和高效改写；
第二个是对一些核心算子进行了高效的优化和改写。以 TensorFlow 中的 CPU Matmul 模式为例，我们采用 Onnx 的 MLAS 高效计算库进行了替换和改写，从而在核心场景中取得了较大的提升。

3. 中后端优化

中后端的优化，在调研了业界的 TensorRT、TVM 和 XLA 之后，我们选择了与当前 TensorFlow 最贴切的 XLA 方向。我们主要解决了两个问题，首先是如何将 XLA 应用于我们的 inference 场景，因为它的输入 tensor 形状是固定的。因此，我们对之前异构推理引擎中的 Auto Batching 技术进行了升级，开发了一种名为 Static Batching 的技术。这种技术可以将一个相邻时间窗口内的多个请求固定在一个比较规整的 batch size 上。通过这种方式，使 XLA 的输入成为一个固定的输入。

第二个问题是关于 JIT 稳定性的问题。JIT 是 XLA 目前主要使用的编译工具，但它在处理线上稳定性和多异构机型时存在一些挑战，例如无法进行模型的热更新。因此，我们最终选择使用 AOT 方案。整个编译流程图如上图所示：首先进行离线模型训练，然后编译出针对线上特定 GPU 机型的可执行文件，并将其分发到对应机型上。

4. 项目收益

利用 AI 编译技术，在核心训练场景，可以将整个计算图的总数降低一个数量级，并使 GPU 计算效率得到了大幅提升，推理场景的 GPU 计算效率能提升 30%+。在核心 CPU 推理场景中，同延迟下吞吐性能提升了 60%+。

最后是我们看未来要做哪些事情。从业务的角度来看，未来我们将面临几个挑战。首先，DAU 将更高；其次，模型将更加复杂；最后，应用场景将更加多样化。因此，我们需要具备更强的模型工程能力，以应对更多的模型需求、更高的训练吞吐量和更大的模型规模，以及更强大的计算能力。

05 展望未来

从业务角度来看，未来将面临三大挑战，即更高的 DAU、更复杂的模型，以及更多的场景。因此我们的训练推理引擎需要更多的模型、更高的训练吞吐、更大的模型和更旺盛的算力。

具体工作主要有四个方向，首先是 AI 编译技术，我们需要不断迭代；第二是 HPC 同步训练，以提供更高的算力和更高的训练吞吐量；此外，我们还需要异构参数服务器和异构弹性训练引擎。

最后也欢迎志同道合的朋友加入我们的团队。

以上就是本次分享的内容，谢谢大家。

分享嘉宾

INTRODUCTION

曾鸣堃

小红书

训推计算引擎负责人

小红书训推计算引擎负责人，17 年毕业于中国科学技术大学，目前负责搜广推场景下的模型训练和推理引擎侧研发工作。

活动推荐

往期推荐

基于 tugraph-analytics 的实时业务数据异常归因诊断

大语言模型在图推荐系统中的融合与优化策略

Data+LLM：金融真实场景的技术创新实践

京东广告稀疏大模型训练与推理 GPU 优化实践

好的数据治理怎么做？

销售易基于 Lakehouse 的实时分析提升用户数据体验实践分享

Velox内存管理深度解析：从基础到高级特性

Apache Hudi 从零到一：全面解读写入索引（四）

Apache Hudi 从零到一：理解写入流程和操作（三）

用最酷的RAG，训最猛的大模型！

点个在看你最好看

SPRING HAS ARRIVED

原创精选813

小红书8

原创精选 · 目录

上一篇基于 tugraph-analytics 的实时业务数据异常归因诊断下一篇企查查的数据降本增效之路

继续滑动看下一个