深度优化 TorchRec：提升 PyTorch 推荐系统性能

本文深入分享了针对 PyTorch 官方推荐系统库 TorchRec 的性能优化工作。TorchRec 是 PyTorch 生态中用于构建大规模推荐系统的核心组件，支持在 PyTorch 框架上进行高效的 Embedding 训练与管理。本次优化的目标非常明确：首先，针对 MLPerf 基准测试中的 DLRM 模型进行专项优化，该场景涉及 16 个 DGX 节点，对扩容性要求极高；其次，确保优化过程不破坏 TorchRec 现有的 API 接口，保持向后兼容性；再次，尽量不改动其底层架构，倾向于在高层级进行性能调优；最后，部分优化如 CUDA Graph 的使用，对模型结构本身有一定约束条件。

TorchRec 整体架构解析

TorchRec 的整体架构设计清晰，主要包含三层结构。最上层是 TorchRec 的 API 层，提供简单易用的封装接口，用户可以利用这些 API 快速配置不同的 Embedding 策略，实现数据分片（Sharding）以及在训练过程中构建流水线（Pipeline）。中间一层是 TorchRec 内部的 Module 层，包含了具体的 Sharding 实现逻辑，将 Embedding 功能划分为 BaseSparseFeatureDist、BaseEmbeddingDist 和 BaseEmbeddingLookup 三个核心部分。这一层主要由 Python 代码构成，基于 PyTorch 的 nn.Module 机制实现。最底层则是 FBGM（Facebook Graph Machine Learning），这是一个 C++ 层面的库，其中包含了在 GPU 上高效实现的稀疏相关算子，为上层提供了高性能的计算基础。

典型使用流程

用户使用 TorchRec 通常遵循三个标准步骤。第一步是使用 EmbeddingBagConfig 和 EmbeddingBagCollection 这两个核心 API 来构建 Embedding 配置，定义特征与 Embedding 表的映射关系。第二步，使用 DistributedModelParallel 完成模型并行（Model Parallel）的分片处理。TorchRec 默认会在 Embedding 层面进行模型并行，将巨大的 Embedding 表切分到每一个 GPU 上，同时也提供了定制化开发的接口，允许用户对特定部分进行精细化的模型并行控制。第三步，使用 TrainPipelineSparseDist API 构建一个完整的训练流水线。在这个流水线中，可以进行流水线调度以及预取（Prefetch）操作，从而显著提升整体性能。由于我们的优化目标是高层级的性能提升，因此重点主要集中在 TrainPipelineSparseDist API 的实现细节中。

优化成果展示

经过一系列优化措施后，我们在 MLPerf DLRM-DCNv2 基准测试中取得了显著成果。测试环境为 16 个 DGX H100 节点。数据显示，优化前 TorchRec 每个迭代（Iteration）耗时约为 7.6 毫秒，优化后降至 3.4 毫秒左右，实现了 2.25 倍的加速比。当前 MLPerf 的世界纪录为 2.3 毫秒，我们的优化方案已非常接近这一极限水平。接下来将详细拆解这一优化方案的具体实施路径。

训练时间线分解

为了定位性能瓶颈，我们将 TorchRec 构建出的 DLRM 整体训练时间线进行了详细分解，总共包括五个关键部分。第一部分为 Input Feature Dist（输入特征分发）。当输入特征从数据库或硬盘读取到 GPU 显存后，需要将其从数据并行（Data Parallel）模式转换为模型并行（Model Parallel）模式，以便后续进行 Embedding 处理。第二部分是在获得 Model Parallel 的特征之后，执行 Embedding Forward。这一步在当前 GPU 上进行 Embedding 表的查询，并执行 Embedding All-to-All 通信。经过 All-to-All 操作后，每个 GPU 能够获取到数据并行的 Embedding 向量，方便后续进行数据并行的 MLP 计算。第三部分是 MLP 的前向传播（Forward）、反向传播（Backward）以及梯度 Allreduce 操作，包括底部 MLP（BMLP）和顶部 MLP（TMLP）。第四部分在完成 MLP 部分的计算后，需要进行 Embedding Backward，这是 Embedding Forward 的反向过程，即先执行 All-to-All 通信，再进行反向传播，最后更新 Embedding 表。最后是第五部分，MLP 的参数更新。这五个部分之间存在严格的依赖关系，②/③/④/⑤在开启流水线后可以顺序执行，它们之间相互依赖，必须按序运行。但是①与其他部分没有直接依赖关系，因此①可以独立进入流水线，例如当前迭代的①可以与下一个迭代的②/③/④重叠执行，从而实现流水线的隐藏延迟效果。

核心优化方案详解

本次优化工作主要总结为两大方向：一是优化 CPU 启动延迟（CPU Launch Latency），二是优化 Input Dist 部分，去除不必要的中间操作。

1. CUDA Graph 降低 CPU 开销

我们使用了 CUDA Graph 技术将 MLP 和 Allreduce 部分的操作捕获下来。通过创建 CUDA Graph Object，我们可以捕获中间的 Kernel 启动过程。在后续执行时，无需再经历 CPU 的调用序列，可以直接调用 CUDA Graph 对象执行，从而节省了大量的 CPU 时间。使用 CUDA Graph 有一个关键前提，就是必须保证 CUDA Graph 的内存地址不能发生变化。对于 TorchRec 而言，Embedding 部分（包括输入特征分发、Embedding Forward、Embedding Backward）都依赖于动态变化的 Embedding 输入特征数量，因此无法直接应用 CUDA Graph。但对于 MLP 计算部分（即上述时间线中的③），由于其计算图相对固定，是可以对其进行 CUDA Graph 捕获的。这是第一个关键的优化点。