深度优化 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 捕获的。这是第一个关键的优化点。

2. 多线程 Kernel 启动解决阻塞

TorchRec 原生实现在同一时间只使用一个线程来启动 Kernel。Kernel 启动的顺序通常是②③①④⑤。之所以采用这样的顺序，是因为在①阶段中有 3 次 All-to-All 通信操作。我们希望尽早启动①，这样可以让①中的三次通信与③中的计算更好地重叠。然而，这种顺序导致了一个问题：①是在中间阶段启动的，但①中包含的两次 D2H（Device to Host）加同步操作会导致 CPU 等待 GPU 完成数据传输。在这个过程中，由于它是单线程启动，CPU 的时间会被浪费，进而导致④和⑤的 Kernel 启动被阻塞。我们的优化方案是采用一个单独的线程去启动①。这样，即使它在单独线程上因 D2H+Sync 而阻塞，也不会影响 TorchRec 主线程上的④和⑤的 Kernel 启动。这样可以充分利用 D2H+Sync 的时间窗口来并发启动其他 Kernel，最大化硬件利用率。

3. Pinned Memory 优化 D2H Copy

在 PyTorch 本身的实现中，默认使用 Tensor API 进行 D2H 传输。在这个 API 中，默认使用的内存页不是常驻在虚拟内存上面的（Unpinned），这会导致在 CPU 层面产生额外的性能开销，因为操作系统需要频繁地进行页面交换。我们的优化是不直接使用默认的 to_list API，而是手动创建一份 Host 上的 Tensor，这个 Tensor 使用 Pinned Memory（固定内存）。然后先完成 Pinned Memory 的拷贝操作，之后再转换为 List。Pinned Memory 允许 DMA（Direct Memory Access）设备直接访问，减少了 CPU 介入和数据搬运的开销，从而显著减少 CPU 层面的性能损耗。

4. 环境变量减少 NCCL 开销

我们使用了 TORCH_NCCL_AVOID_RECORD_STREAMS 这个环境变量来减少 CPU 上的额外开销。具体使用场景是当 PyTorch 使用多个 Stream 时需要保证内存安全性。例如，如果在当前 Stream 创建了一个张量 a，将其放入 Stream s 中使用，然后删除 a，这时就会引入竞态条件（Race Condition）。因为删除 a 时，它可能正在被 s 异步使用，这会产生安全隐患。PyTorch 本身提供了一个 API 叫做 record_stream，用于保证删除 a 时去做轮询，确保内存只有在被使用完之后才会被释放。虽然 record_stream 提供了安全性保证，但它的工作状态必须不断地去轮询，这会造成一定的性能开销。因此，对于某些安全但并不高效的场景，我们寻求替代方案。另一种方式是当使用 a 时，不要在 s 里面做 record_stream，相反我们做 Stream 之间的同步。在使用完 s 后，使用当前的 Stream 等待 s 完成，就可以保证 delete a 操作不会在 s 被使用结束之前去执行。这样避免了使用 record_stream 时的轮询开销，性能更好。PyTorch 本身底层的 NCCL 库不知道用户的具体使用方式，所以默认进行 record_stream，这会引入性能损失。如果我们像场景二那样使用了 Stream Synchronize，就可以通过将 TORCH_NCCL_AVOID_RECORD_STREAMS 设为 1，去掉默认的 record_stream，并且对内存安全性没有任何负面影响。

5. 优化 Input Dist 减少通信

最后一个优化是去掉了①阶段中的一些中间操作，具体来说就是 Input All-to-All 和 D2H+Sync。在①阶段中有 3 个 All-to-All 操作，分别是 a2a0、a2a1 和 a2a2。其中，a2a0 用来传输 a2a1 和 a2a2 传输需要的元数据。我们发现，这些元数据并不是必须要通过网络通信来传输的，而是可以在本地通过计算得到的。因此，我们可以通过重计算（Recomputation）的方式，把 3 次 All-to-All 合并成 2 次，节省一次昂贵的网络通信。此外，我们发现 a2a0 传递的信息，只会因为每个 GPU 上的 Batch Size 变化而产生改变。如果 Batch Size 没有变化，就没有必要重复做 a2a0 的同步，结果是不会改变的。因此，我们的优化策略是把 a2a0 和相关的 D2H+Sync 去掉。当 Batch Size 没有变化时，尝试复用上一个 Iteration 的结果，只有在 Batch Size 发生变化的特殊情况下，才会重新执行这个 a2a0 操作。在训练的大部分时间内，Batch Size 都是固定的，因此可以从中获得持续的性能收益。

性能指标分析

最后展示一下各项优化带来的具体性能结果。我们挑选了 3 个对性能影响最大的优化项进行对比展示，包括 CUDA Graph、去掉一次 All-to-All 通信，以及多线程的 Kernel 启动。这里展示了 5 个关键指标，包括 GPU 上每一个 Iteration 的训练时间，即 CPU 上 Iteration 从开始到结束的时间，去掉了 CPU 同步 GPU 的时间，纯粹反映 CPU 上的时间开销。然后是 Exposed Input Dist，用来描述 Input Dist Prefetch 暴露给 CPU 的时间，最好的情况是可以完全隐藏掉这部分开销。将这些优化拆开来看，所有的优化大致上都可以减少在 CPU 上的 Overhead，同时能够减少在 GPU 端对端的训练总时间。最主要的收益来自于 Input Dist 部分的优化，我们的改进使 Input Dist 部分被更好地隐藏在流水线之下，减少了 CPU 等待时间。

总结与适用性

这些优化主要针对 Embedding 本身的一些通用特性，可适用于大部分推荐系统场景。CUDA Graph 部分与网络结构有关，它不会依赖动态的输入，如果是像 Transformer 等依赖动态输入的情况，则需要使用其它方法去做优化。总体而言，通过上述多层面的系统级优化，我们成功在保持 TorchRec 原有架构稳定性的前提下，显著提升了大规模推荐系统的训练效率，为工业界落地提供了有力的性能保障。

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