微信基于 PyTorch 的大规模推荐系统训练实践 _PyTorch

文章插图
本文将介绍微信基于 PyTorch 进行的大规模推荐系统训练。推荐系统和其它一些深度学习领域不同，仍在使用 Tensorflow 作为训练框架，被广大开发者诟病。虽然也有使用 PyTorch 进行推荐训练的一些实践，但规模较小，也没有实际的业务验证，很难推动业务尝鲜。
2022 年 2 月，PyTorch 团队推出了官方推荐库 TorchRec 。我们团队在 5 月开始在内部业务上尝试 TorchRec，并且与 TorchRec 团队展开了一系列的合作。在几个月的试用过程中，我们体会到 TorchRec 非常多的优点，也感受到 TorchRec 在超大规模模型上仍存在一些不足。针对这些不足，我们设计了扩展功能来填补它的问题。在 2022 年 9 月，我们设计的扩展功能 dynamic embedding 已经正式合进了 TorchRec 的主分支，目前仍在与官方团队持续优化。

文章插图

一、TorchRec 可以为我们带来什么

文章插图

我们先来聊聊 TorchRec 可以给我们带来什么？我们都知道推荐系统往往和公司的现金流直接挂钩，试错成本非常高，所以大家需要的是一个经过了业务测试的框架。这也是为什么之前的一些基于 PyTorch 的推荐框架都未曾被广泛应用过。而 TorchRec 作为一个官方的推荐框架，在 2022 年 1 月份推出之时，Meta就已经利用它在 Instagram Reels 业务上顺利训练并上线了一个 1250 亿参数的模型，成为了一个经过业务测试的 PyTorch 框架。有了 Instagram 这样一个大业务的支撑，让我们有了更多信心，终于可以去理性地考量一个基于 PyTorch 的推荐框架有什么样的优势了。

文章插图

对于团队中的不同成员，TorchRec 有不同的好处。首先，对于团队中占绝大多数的算法工程师而言，PyTorch 推荐框架让大家终于可以享受到像 CV、NLP 工程师体会到的那种更人性化的动态图和调试的体验。
另外，PyTorch 极好的兼容性——一个基于 PyTorch1.8 做的模型，不需要改一行代码就可以在最新版本 1.13 上运行——让算法工程师终于可以放心地升级框架，从而享受到最新的框架功能和更优秀的性能。而反观一些基于 TensorFlow 的推荐框架，往往被卡在 TensorFlow 的某一个版本上，例如很多团队可能还在使用基于 TensorFlow 1.x 的内部框架。TensorFlow 1.x 在 2021 年 1 月份就已经停止维护了，这就意味着在近两年的时间内，所有新出的 bug、新出的特性都无法得到很好的支持。使用过程中遇到的问题，也只能靠内部维护团队去修复，增加了额外的成本。及时的框架升级还可以带来免费的速度提升，高版本的 PyTorch 往往匹配更高版本的 CUDA，以及像 CUDA graph 等的一些新特性，可以进一步提升训练速度，提升训练效率。

文章插图

除了算法工程师，框架团队也是推荐团队的重要组成部分。公司中的框架团队会在选取开源框架之后基于内部需求进行二次开发。对于他们来说，一个 PyTorch 的推荐框架会带来更简化的开发体验。很多传统的 TensorFlow 推荐框架会模仿 TF serving 来做一个基于 C++ session 的扩展——这样的设计方案在当时算是非常先进的方案——但这使得只改一行代码也需要完整地编译整个 TensorFlow，耗时很长，甚至还要在解决在内网下载外部的依赖之类的琐碎问题，开发体验不太好。
使用 PyTorch 不会遇到这样的问题，因为 PyTorch 以 Python/ target=_blank class=infotextkey>Python 哲学为核心，它希望大家能够自如地进行扩展。我们在进行二次开发的时候，只需要用 pybind11 这样比较成熟的 Python 库封装一下，把我们的库打包成一个动态链接库，就可以加载了。这样自然整体编译速度会快很多，同时学习成本也会低不少。
前面提到 PyTorch 是一个向后兼容性非常好的框架，这让维护团队不需要去维护多个版本，很多共性的问题都可以得到官方的解决，大家就可以专注于特化需求，团队人员效率就会有明显提升。
上面介绍的都是 TorchRec 作为一个 PyTorch 推荐框架的优势，让我们感到非常开心的是，TorchRec 团队没有止步于做一个 PyTorch 推荐框架。他们观察了现有推荐模型以及硬件的特点，在框架中加入了许多的新特性，使得 TorchRec 相比于传统的推荐框架有明显的性能优势。接下来我会选择其中的几个来进行介绍，分别是 GPU embedding，TorchRec 里面优秀的 GPU kernel，还有 TorchRec 能够根据网络通信进行的 embedding 划分。