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

导读 本文将分享对 PyTorch 上的推荐系统库 TorchRec 的优化工作。

TorchRec 是 PyTorch 官方开发的，支持在 PyTorch 上做大规模 embedding 的一个推荐系统训练库。优化目标包括：

• 首先，是针对 MLPerf 中 DLRM 本身的 benchmark，有 16 个 DGX nodes，是扩容性非常高的一个场景；
• 另外，我们希望不破坏 TorchRec 本身的 API；
• 并且不破坏其本身的架构，更倾向于做一些 high level 的优化；
• 其中有一些优化跟模型结构相关，比如后文中会介绍的 cuda graph，对模型结构本身有一定要求。

上图展示了 TorchRec 的整体架构，包括三层。最上层是 TorchRec 的 API，主要提供一些简单易用的 wrapper，用户可以很轻松地利用这些 API 来配置不同的 embedding，去做 sharding，在 training 当中做流水线的工作。中间一层是 TorchRec 内部 module，包括一些 sharding 的实现，把 embedding 分成了 3 个部分，BaseSparseFeatureDist、BaseEmbeddingDist 和 BaseEmbeddingLookup。中间这一层主要都是 python 上的代码，就是 PyTorch 当中的 nn.module。最底层是 FBGM，是 C++ 层面的 library。其中包括 GPU 上高效实现的 sparse 相关的 operator。

上图中展示了一个用户使用 TorchRec 的例子。可以看到，使用 TorchRec 包括三步：

第一步是使用 EmbeddingBagConfig 和 EmbeddingBagCollection 这两个 API 构建 embedding。

第二步，使用 DistributedModelParallel 完成 sharding。TorchRec 本身默认在 embedding 上做 model parallel，会把 embedding 切分到每一个 GPU。同时也提供了支持定制化开发的接口，支持对部分内容做 model parallel。

第三步，使用 TrainPipelineSparseDist API 构建一个训练流水线，在流水线当中可以做 pipeline，以及 prefetch，这样会有更好的性能。因为我们更希望 high level 的优化，所以优化主要集中在 TrainPipelineSparseDist API 中。

上图是优化之后的结果。开头提到，我们是针对 MLPerf DLRM-DCNv2 在 16 DGX H100 的节点上。图中是每一个 iteration 的耗时，可以看到原本 TorchRec 每一个 iteration 需要耗时 7.6 毫秒，优化之后为 3.4 毫秒左右，有 2.25 倍的加速。当前 MLPerf 的世界纪录是 2.3 毫秒。

接下来将详细介绍这一优化方案。

上图对 TorchRec 构建出来 DLRM，整体的训练 timeline 进行了分解，总共包括 5 个部分，如图中圆圈所标注出的。

• 第一部分为 input feature dist。当 input feature 从数据库、硬盘读到 GPU 之后，需要把它从 data parallel 转换成 model parallel，这样才可以在后面进行 embedding 处理。
• 第二步是在有了 model parallel 的 feature 之后，进行 embedding forward。Embedding forward 是在当前 GPU 上进行 embedding table 的查询，以及 embedding all-to-all。经过 all-to-all 之后，每一个 GPU 上能够拿到 data parallel 的 embedding vector，这样可以方便后面做数据并行的 MLP。
• 第三部分是 MLP 的 forward 和 backward，以及 gradient allreduce。如图中标记的 BMLP（bottom MLP）和 TMLP（top MLP）。
• 做完 MLP 部分的 forward、backward、allreduce 之后，需要做 embedding backward，即第四部分，可以看到它是 embedding forward 部分的反向，也就是先去做 all-to-all，再做 backward，然后做 embedding table 的更新。
• 最后是第五部分，MLP 的 update。

这五个部分之间的依赖关系为，②/③/④/⑤在开启 pipeline 之后是可以顺序执行的。他们之间有相互依赖关系，必须顺序执行。但是①和其他没有依赖关系，所以①可以做 pipeline，比如当前的 iteration，可以做下一个 iteration 的①，如图中所展示的，①和②/③/④是重叠的关系。以上就是整体的时间线的分解。

优化可以总结为两大部分，第一大部分是优化 CPU launch latency。我们使用了 cuda graph 把 MLP 和 allreduce 部分做 cuda graph，然后使用 multi-threading kernel launch 解决中间部分的 launch bond。我们还优化了其 D2H 的 copy 部分，使用 pinned memory copy 减少 CPU 上的 overhead。并且使用了 TORCH_NCCL_AVOID_RECORD_STREAMS 这个环境变量，后面会介绍它是如何帮我们减少在 patch 里面使用 NCCL 时造成的开销的。

第二大部分是优化它本身的 input dist 部分，去掉第 0 个 operator 和相关的 D2H+sync。

接下来详细介绍这些优化。

首先是 cuda graph 的部分。如图中绿色标记的部分，我们使用 cuda graph object 捕获中间的 launch kernel，后续再执行就不需要再经历 CPU 的调用，可以直接调用 cuda graph，这样可以节省在 CPU 上的时间。Cuda graph 使用有一个条件，就是必须保证 cuda graph 的内存地址不能够变化。对于我们来说，embedding 部分包括①/②/④，它们都需要依赖动态的 embedding 输入 feature 的个数，所以没有办法对其进行 cuda graph，但对于③这个部分，是可以对它 cuda graph 的。这是第一个优化。

第二个优化是多线程的 kernel launch。TorchRec 本身在同一个时间只会使用一个 thread to launch kernel。Launch kernel 的顺序是②③①④⑤。之所以是这样的顺序，是因为在①当中有 3 次 a2a，3 次 communication。我们希望尽早把①launch 起来，这样就可以更好地让①当中的这三个 a2a 的通信，和③中的计算更好的 overlap。但这样会导致一个问题，就是①是在中间做 launch 的，但是①当中两次 D2H+sync 的操作会导致 CPU 等待 GPU，在这个过程中，因为它是单线程的 launch，所以就会导致 CPU 的时间被浪费，使得④和⑤的 kernel launch 阻塞住。我们的优化是采用一个单独的 thread 去 launch①。这样它在单独的 thread 上阻塞了，阻塞的同时可以在 TorchRec 本身的 thread 上 launch④和⑤。这样可以充分利用 D2H+sync 时间来 launch kernel。这是第二个优化。

第三个优化是使用 pinned memory 完成 D2H copy。在 PyTorch 本身的实现当中，使用一个 tensor API 进行 D2H。在这个 API 当中，默认的实现是 Mpinned memory，内存页不是常驻在虚拟内存上面的，这样会导致在 CPU 层面有一些性能上的开销。我们做的优化，是不去直接使用 to list API，而是去创建一份 host 上的 tensor，这个 tensor 是 pinned memory，然后先完成 pinned memory copy 操作，之后再做 to list，这样可以减少 CPU 上的性能开销。

第四个优化是使用环境变量 TORCH_NCCL_AVOID_RECORD_STREAMS 减少 CPU 上的开销，具体使用场景就是 PyTorch 在使多个 stream 时需要保证内存安全性。比如图中最左边的场景，首先在当前这个 stream 创建了一个 a，然后把这个 a 放在这个 stream s 里面使用，然后再 delete 这个 a，这时就会引入 race condition。因为删除 a 时，它可能正在被 s 使用，这样的会产生不安全的隐患。

PyTorch 本身提供了一个 API，叫做 record_stream。在这个场景中表示的是，record_stream 可以保证 delete 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，并且对安全性也没有任何影响。

最后一个优化是去掉了①中的一些中间操作，具体来说就是 Input a2a 和 D2H+sync。

可以看到在 1 当中有 3 个 a2a，分别是 a2a0、a2a1 和 a2a2，其中，a2a0 用来传输 a2a1 和 a2a 传输需要的元数据，但是我们发现，这些元数据并不是必须要通过通信来传输的，而是可以在本地通过计算得到的，所以我们可以通过重计算，把 3 次 a2a 合并成 2 次，节省一次 a2a。

我们发现 a2a0 传递的信息，只会因为每个 GPU 上的 batch size 变化而产生改变。那么如果 batch size 没有变化的话，就没有必要做 a2a0 sync，本身结果是不会改变的。所以我们的优化就是把 0 和相关的 D2H +sync 去掉。当 batch size 没有变化时尝试去复用上一个 iteration 结果，只有在 batch size 发生变化时，才会做这个 a2a0。在训练的大部分时间内，都是不会变化的，因此可以从中获得收益。

以上就是对主要优化的介绍。

最后展示一下性能上的结果。挑选了 3 个性能影响比较大的优化来展示其效果，包括 cuda graph，去掉一次 a2a，以及多线程的 launch。这里展示了 5 个指标，包括 GPU 上每一个 iteration 的训练时间，即 CPU 上 iteration 从开始到结束，去掉 CPU 同步 GPU 的时间，就是 CPU 上的时间开销。然后是 exposed input dist，用来描述 input dist prefetch 暴露的时间，最好的情况是可以完全隐藏掉。

我们把这些优化拆开之后，可以看到所有的优化大致上都可以减少在 CPU 上的 overhead，同时能够减少在 GPU 上端对端的训练的时间。最主要的收益来自于 input dist，我们的优化使 input dist 部分被隐藏得更好。

这些优化都是对 embedding 本身的一些优化，可适用于大部分情况。Cuda graph 部分与 network 结构有关，它不会依赖动态的输入，如果是像 transformer 等依赖动态输入的情况，则需要使用其它方法去做优化。

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

​

如何学习大模型

现在社会上大模型越来越普及了，已经有很多人都想往这里面扎，但是却找不到适合的方法去学习。

作为一名资深码农，初入大模型时也吃了很多亏，踩了无数坑。现在我想把我的经验和知识分享给你们，帮助你们学习AI大模型，能够解决你们学习中的困难。

我已将重要的AI大模型资料包括市面上AI大模型各大白皮书、AGI大模型系统学习路线、AI大模型视频教程、实战学习，等录播视频免费分享出来，需要的小伙伴可以扫取。

一、AGI大模型系统学习路线

很多人学习大模型的时候没有方向，东学一点西学一点，像只无头苍蝇乱撞，我下面分享的这个学习路线希望能够帮助到你们学习AI大模型。

二、AI大模型视频教程

三、AI大模型各大学习书籍

四、AI大模型各大场景实战案例

五、结束语

学习AI大模型是当前科技发展的趋势，它不仅能够为我们提供更多的机会和挑战，还能够让我们更好地理解和应用人工智能技术。通过学习AI大模型，我们可以深入了解深度学习、神经网络等核心概念，并将其应用于自然语言处理、计算机视觉、语音识别等领域。同时，掌握AI大模型还能够为我们的职业发展增添竞争力，成为未来技术领域的领导者。

再者，学习AI大模型也能为我们自己创造更多的价值，提供更多的岗位以及副业创收，让自己的生活更上一层楼。

因此，学习AI大模型是一项有前景且值得投入的时间和精力的重要选择。