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

文章预览

导读   本文将分享对 PyTorch 上的推荐系统库 TorchRec 的优化工作。 分享嘉宾｜ 刘仕杰 NVIDIA 加速计算专家 编辑整理｜ 王跃辉 内容校对｜李瑶 出品社区｜ DataFun TorchRec 是 PyTorch 官方开发的，支持在 PyTorch 上做大规模 embedding 的一个推荐系统训练库。优化目标包括： 首先，是针对 MLPerf 中 DLRM 本身的 benchmark，有 16 个 DGX nodes，是扩容性非常高的一个场景； 另外，我们希望不破坏 TorchRec 本身的 API； 并且不破坏其本身的架构，更倾向于做一些 high level 的优化； 其中有一些优化跟模型结构相关，比如后文中会介绍的 cuda graph，对模型结构本身有一定要求。 上图展示了 TorchRec 的整体架构，包括三层。最上层是 TorchRec 的 API，主要提供一些简单易用的 wrapper，用户可以很轻松地利用这些 API 来配置不同的 embedding，去做 sharding，在 training 当中做流水线的工作 ………………………………