使用 PyTorch 构建高效推荐系统推理系统

为什么选择 PyTorch 构建推荐系统

PyTorch 已成为 AI 社区事实上的框架，绝大多数前沿研究，尤其是在推荐系统、检索和排名等领域，都是使用 PyTorch 进行的。开发者渴望尽快将最新的模型进展投入生产。基于 PyTorch 的推荐推理系统非常契合这一需求，既能实现（1）高效率，又能（2）在生产环境中快速采用新模型。

在本篇博客中，我们将讨论使用 PyTorch 构建高性能推荐推理系统的设计。基于这些设计原则的方法已得到充分验证，并成功处理了极高流量，展现了强大的效率和可靠性。我们基于 PyTorch 的推荐推理系统是 Meta 最关键机器学习工作负载的基石。该系统为包括信息流（Feed）、广告（Ads）、Instagram、Reels、Stories 和 Marketplace 在内的全球业务提供支持，管理着多种多样的 ML 架构。从基础的深度学习推荐模型 (DLRM) 的复杂扩展，到诸如 DHEN (深度分层集成网络)、HSTU (分层序列转换单元)、Wukong 等前沿创新建模技术，该系统均能游刃有余。

典型的推荐模型从研究到生产的推理工作流

整体工作流

训练完成后，模型定义及其训练权重会被交付用于推理，从而在训练阶段和推理阶段之间建立明确的契约。然而，直接在生产推理环境中使用训练模型运行是非常低效的，且无法满足现实应用的性能要求。

为了解决这个问题，我们需要快速且可靠地将训练好的模型发布到生产环境，同时支持随着模型改进或再训练而进行的频繁更新。这种模型众多、版本各异的动态环境，要求必须有一个强大的转换流水线，将训练模型转换为优化的推理模型。该流水线确保生成的推理模型文件是针对硬件利用率进行优化的，从而实现高吞吐量（QPS，即每秒查询次数）并满足严格的延迟要求。简而言之，一个专门用于将训练模型转换为生产就绪推理模型的系统，对于维护模型部署过程的敏捷性、可扩展性和性能至关重要。

从训练模型到生产推理的转换流程

定义推理模型与权重映射

训练模型通常包含仅在训练期间必要的组件，例如损失函数和某些正则化技术。最佳实践是定义一个专门的推理模型，它在镜像训练模型的前向逻辑的同时，允许进行仅限推理的优化。此外，必须建立推理模型参数与训练模型权重（检查点）之间的映射，尤其是当完全限定的参数名称在训练和推理之间存在差异时。此映射应在整个推理模型准备过程中进行维护和更新。

从 Python 模型中捕获计算图

为了实现高效推理，必须对推理模型应用一系列模型转换。执行这些优化需要将 Python 定义的 PyTorch 模型转换为图表示。捕获 PyTorch 模型的计算图是一项具有挑战性的任务。使用 torch.fx 提取 FX 图是一种常见的做法。此方法假设模型架构不包含循环结构。对于具有复杂控制流的子模块，可以将它们标记为叶子节点，以简化图提取过程。

最近，torch.export 已成为捕获计算图的更成熟工具，为具有控制流的模型提供了更好的支持。然而，由此生成的 PT2IR（一种专门的 FX 图）可能非常底层且解耦，这可能会使某些模型转换变得复杂。

模型转换与优化

捕获 FX 图后，可以通过模型转换传递应用各种优化。以下是一些常见的转换操作：

模型拆分 (Model Splitting)：对于分布式推理场景，通常需要将完整的前向“forward”图拆分为更小的子图。每个子图代表一个子模块的前向传播，从而实现跨多个设备或主机的分布式执行。此外，这些转换可以将相似的计算组合在一起，进一步提高整体效率。
算子融合 (Operator Fusion)：可以用单一的融合实现来替换多个操作，从而提高效率。这可以通过交换子模块或应用图级转换来实现。
量化 (Quantization)：类似于算子融合，某些层（例如线性层）可以用量化版本替换，以减少内存使用并提高推理速度。TorchAO 提供了对支持 PT2 的线性量化的支持。
编译 (Compilation)（又称降低/Lowering）：模型编译技术通常作为转换过程的一部分提前应用。此步骤将模型代码转换为更适合目标推理设备的底层表示。（详情请参阅下文的 AI 编译器部分。）

图转换示例：从完整的前向图到拆分图

模型序列化

标准 PyTorch 模型使用 pickle 格式进行存储，但由于向后兼容性差和 Python 依赖问题，这种方法不足以满足生产需求。为了应对这些挑战，可以使用以下几种序列化方案：

解决方案	描述	优点	缺点
TorchScript	通过脚本化或跟踪捕获 TorchScript IR，并保存为 TorchScript 格式。	1) 成熟且具有强大的向后兼容性支持 2) 可靠的控制流支持	1) 模型定义存在一些限制（例如不支持复杂数据结构） 2) 已被弃用，不再受支持
torch.export	将 PyTorch 模型导出为 PT2IR。	1) PT2 中官方推荐的序列化模型方式 2) 正在积极开发中	1) 控制流可能需要额外处理
torch.package	直接将相关的 Python 模块导出为源代码和 pickle 对象。	1) 极大的灵活性	1) 可能需要手动定义模块边界 2) 依赖于 Python

无论使用哪种序列化格式，生成的制品都应该是 zip 文件。这允许通过解压缩内容进行轻松检查和调试。处理后的权重也可以打包在 zip 文件中。我们优先考虑使用 torch.export 进行新模型开发，而不是像 TorchScript 和 torch.package 这样的旧工具。随着 TorchScript 被弃用，torch.export 提供了一条更稳健的前进道路，并具有积极的功能开发；同时，通过支持独立于 Python 的运行时，它还提供了优于 torch.package 的必要性能。

模型加载与执行

一旦准备好推理模型，您将获得一组推理模型文件。对于超大模型，可能需要分别加载模型结构和权重，这可能需要自定义保存和加载逻辑。

加载模型文件后，运行时开始处理推理请求。由于 PyTorch 除了模型执行外并未原生提供服务功能，因此需要额外的服务器层来管理推理服务。以下概述了高效且可扩展的推荐系统 PyTorch 推理服务器的关键特征：

轻量级 PyTorch 执行器封装 (Executor Wrapper)

服务器将请求转换为 PyTorch 模型输入。该封装器应尽可能轻量，以确保效率。

高效且灵活的 API

在分布式环境中，模型的不同组件通过 API 进行通信，这就需要精确的语义定义，例如指定批次维度和其他相关参数。
基于 Tensor 的 API 与 PyTorch 模型的前向方法非常契合。
零拷贝 (Zero-copy) API 允许我们原地更新模型，从而在无需额外容量加载两个模型版本的情况下，高效且无缝地从一个模型版本过渡到下一个版本。

DAG 表示与执行器

具有相似特征的模块（例如所有 embedding bags）可以分组到专门的子模块中进行批量执行。
模型拆分后，原始的前向函数被表示为有向无环图 (DAG)，每个节点对应一个子模块。需要一个执行器来管理该 DAG 的执行。
DAG 节点可能部署在多个主机上，这要求支持远程执行。在这种情况下，高效的通信库对于确保分布式组件之间无缝且高性能的交互至关重要。

优化

在上一节中，我们概述了使用 PyTorch 构建稳健、高效且可扩展的推荐推理系统的核心原则，该系统能够处理高流量并满足严格的生产要求。为了进一步提升系统性能，我们将在下面讨论几个关键的优化策略。

GPU (加速器) 推理

随着新模型架构的出现，计算需求显着增加。CPU 通常难以满足在线运行此类模型的延迟要求，这使得 GPU 等加速器成为自然之选。然而，在单个 GPU 上运行整个模型可能效率低下，且模型可能无法容纳在单个设备的内存限制内。因此，将模型拆分为多个段并在 GPU 上执行计算最密集的层是一种实用的方法。

此外，GPU 内核启动开销可能很大。为了缓解这种情况，将请求进行批处理可以减少内核启动次数并提高整体吞吐量。

C++ 运行时

虽然运行 PyTorch 模型最直接的方法是通过 Python，但 Python 运行时会带来明显的开销，特别是随着 QPS（每秒查询次数）的增加。通常，当 QPS ≥ 100 时，Python 的开销变得显着，而在 QPS ≥ 1000 时，它可能成为严重的瓶颈。

对于高 QPS 场景（每主机 ≥ 100），我们建议使用 C++（或 Rust）运行时。TorchScript（针对 TorchScript 模型）和 ExecuTorch（针对使用 torch.export 保存的模型）都提供了 C++ 运行时。最近，开发重点转向了一个新的运行时 torch.nativert，旨在跨服务器执行 torch.export 模型，以此替代在上一届 PyTorch 大会上已被弃用的 TorchScript 运行时。

分布式推理 (DI)

将整个推理模型作为一个整体运行可能是低效甚至不可行的。相反，将模型拆分为多个组件并将其分发到不同的工作节点上，既能提高效率，又能扩展到更大的模型规模。常见的 DI 模式包括：

CPU-GPU DI：将输入处理和轻量级计算分配给 CPU，同时将模型的计算密集型层卸载到 GPU。
Embedding-Dense DI：将嵌入表（Embedding tables）分组到专门的子模块中，这些子模块可以在单独的主机上提供服务（类似于传统的参数服务器）。稠密层（Dense layers）虽然较小但计算密集，可以进行分组并一起执行以提高效率。
稠密模型并行 (Dense Model Parallelism)：将单个稠密网络拆分为多个子网络，这些子网络可以并行执行（无论是在同一设备的通过不同 CUDA 流，还是跨多个设备），从而实现选择性降低和并行执行。

AI 编译器和高性能算子库

为了获得最高性能，开发者可能想用 C++/CUDA 重写模型定义并直接运行。然而，这种方法扩展性不佳。相反，AI 编译器可以自动化此过程，生成高度优化的制品。可选方案包括：

AOTInductor (torch.compile)
AITemplate
TensorRT

这些编译器生成新的编译后制品，并与序列化模型一起打包。对于生产环境的推荐系统部署，出于性能原因，首选 C++ 运行时。这排除了使用依赖 Python 的 JIT 工作流（如 torch.compile）；相反，我们使用 Ahead-of-Time (AOT) Inductor 将模型预编译为可部署在 C++ 中的静态运行时制品。

AI 编译器利用高性能算子库来最大化各种硬件平台上的计算效率，包括：

CUTLASS/CuTeDSL 用于 Nvidia GPU
Composable Kernels / AITER 用于 AMD GPU
ZenDNN 用于 AMD CPU
OneDNN 用于 Intel CPU

请求合并 (Request Coalescing)

为了最大化效率，应该将请求合并（批处理）在一起。这需要了解每个输入的语义，特别是哪个维度代表动态批大小，以便适当地连接请求。模型的 forward 方法应标记批次信息以促进合并，并且运行时必须支持此功能。

表批处理嵌入 (Table Batched Embedding)

在 PyTorch 中查询嵌入表会产生巨大的算子内核启动开销，特别是在处理数十、数百甚至数千个表时。由于嵌入查找是数据传输密集型操作（类似于哈希映射查询），将嵌入表（embedding bags）进行批处理并在单次调用中查询所有表，可以极大地降低开销并提高数据传输效率。

量化

模型的嵌入层和稠密层都可以从量化中受益匪浅：

嵌入 (Embeddings)：bf16 和 int8 等数据类型通常是安全的，int4 通常也是可以接受的。不同的表和行可能具有不同的数值敏感度。PyTorch 支持按表量化，即使对于表批处理嵌入也是如此，允许开发者自定义量化策略。一些表甚至可以使用 int1 或 int2 配置。
稠密层 (Dense Layers)：稠密层对量化更敏感。通常，fp16 和 bf16 对于整个稠密子模块是可以接受的，但也存在例外，例如 fp16 可能缺乏足够的范围，而 bf16 可能无法提供足够的精度。为了进一步提高效率，可以在层级应用 fp8 和 fp4，尽管这通常需要手动调优。

所有量化策略都应通过精度评估进行验证。TorchAO 提供了对线性层和卷积层的支持，建议从此入手。

增量更新 (Delta Update)

模型新鲜度对于服务推荐模型至关重要。随着模型变得越来越大，加载整个模型变得越来越昂贵。一种平衡的方法是应用部分权重更新（增量更新）。虽然实现数据传输协议很简单，但调整权重加载步调对于避免扰动服务至关重要。嵌入表通常对部分更新的容忍度更高，而稠密模块则更敏感。对于稠密模块，我们建议使用缓冲区模块来支持完整模块交换，而不是单独更新权重。

开发者体验

Python 运行时

为了简化推理流程的开发和调试，我们建议提供一个轻量级的 Python 运行时环境（相对于使用 C++ 运行时）。这种方法允许开发者有效地确定问题是源于运行时还是模型本身。此外，它还简化了为调试目的添加指令的过程。

随着自由线程 Python 的引入，Python 生态系统内的运行时和通信开销都可以进一步降低。这一进展也使得在生产环境中部署 Python 运行时变得越来越实用。

基于模块交换的转换

历史上，基于图的转换对于模型作者来说很难理解和调试，这主要是由于图操作的复杂性以及原始堆栈跟踪信息的丢失。为了解决这些问题，我们建议将此类优化提前到推理模块编写过程的早期。通过采用全面的、原生 PyTorch 模块化的工作流，并利用即时模式（Eager mode）转换，我们发现推理开发体验得到了显着提升。

评估流程

为了确保模型和运行时的质量，我们建议实施以下两个评估流程：

精度验证：将推理模型的质量与训练评估结果进行比较。
性能基准测试：重放类生产流量以评估吞吐量和延迟。

结论

在 Meta，我们开发了一个基于 PyTorch 构建的高效推荐推理系统，这对于将前沿研究转化为生产级服务至关重要。本博客详细介绍了一个稳健的工作流，从训练模型定义及其权重开始，逐步通过必要的推理转换步骤，包括图捕获、模型拆分、优化（融合、量化、编译等），最后是序列化。然后，我们概述了高性能推理服务器的要求，强调了轻量级执行器、灵活的基于 Tensor 的 API 以及基于 DAG 的模型执行模型。最后，我们探讨了对于高 QPS、低延迟性能至关重要的先进优化技术，例如利用 GPU/加速器推理、采用 C++ 运行时、实施分布式推理模式、利用 AI 编译器，以及应用诸如请求合并、表批处理嵌入和量化等高级方法。通过遵循这些原则并利用所介绍的开源库，开发者可以构建可扩展、高性能且敏捷的基于 PyTorch 的系统，能够服务于全球要求最苛刻的 ML 推荐工作负载。

使用 PyTorch 构建高效的推荐系统推理系统