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

Column wise 则是从 embedding 长度角度去划分。例如 embedding 总长是128，可以前 64 维和后 64 维放在不同的位置，这样通信会比较均衡，但是在读取的时候，需要和所有的卡或者 PS 通信。

文章插图

划分模式上的差别带来了选型中的 trade-off 。传统的推荐框架会在设计中固定 embedding 的划分方式，而 TorchRec 则在支持了多种划分方式——比如 row wise、column wise，甚至 table wise，data parallel——的基础上，在内部提供了如 Planner、Estimator、PerfModel 等可以根据使用场景的带宽、显存、内存、模型大小等等参数自动地去计算划分的方式的模块。这样就可以根据我们实际硬件情况去最高效地划分 embedding，最高效地利用硬件。这些功能大都是在 Python 里面实现的。方便我们针对内部环境进行客制化，从而不费力地构建出一套最适合于我们内部环境的推荐系统。
二、在百亿模型上的实验效果

文章插图
在我们的实验中，对于 DeepFM、DCN 这样的在标准模，TorchRec 相对于之前的基准的推荐框架会有惊人的 10 至 15 倍的性能提升。拿到了这样的性能收益，让我们有信心把 TorchRec 推到了业务上。

文章插图

对于微信读书精排模型，在对齐精度的基础上，我们发现在真实数据上有 3 倍左右的性能提升，在假数据上甚至有 10 倍左右提升。这里的差别是因为训练读取数据变成瓶颈了，这方面我们还在做进一步的优化。
03
原始方案在千亿及更大模型上的不足

文章插图

前面介绍的基本是百亿级别或者以下的模型，也就是单机就可以放得下的模型。在把 TorchRec 推到更大的模型的时候，我们观察到 TorchRec 的原生设计的一些问题。对于大模型来说，TorchRec 的纯 GPU embedding 方案需要更多的卡——可能原本 8 张卡的训练速度就可以吃进全部数据，但是我们要用 16 张卡放下 embedding，这使得好不容易提升上去的 GPU 硬件利用效率又被拖了下来。
而且对于大模型的场景，算法团队往往会提出 embedding 的动态增删需求，例如删除一周没有访问过的 ID 。TorchRec 的方案是不支持这样特性的。还有，超大模型的业务一般都会涉及诸多团队，迁移基层框架会遇到很大的阻力。我们需要的支持逐步地渐进迁移，而不能让大家一起放下手头的工作，那样的成本过高，风险太大。

文章插图
根据上述的需求，我们考虑如何去修改 TorchRec，使得它能够适应超大规模模型的场景。我们认为在超大规模训练中，仍然需要支持连接远程的 PS，因为远端 CPU PS 已经非常成熟了，非常容易支持 embedding 的动态增添。同时，对于跨团队的合作，可以用 PS 来隔离开训练和推理，实现渐进的迁移。

文章插图

那么接下来一个问题就是该如何引入 PS 。如果把 PS 直接连到 GPU embedding 上，每个迭代步还是要去访问远端的 PS，会重新使网络和 CPU 整体操作的占比提升，GPU 利用率又被拉下来了。
04
微信团队的 dynamic embedding 如何解决问题

文章插图

这个时候我们发现单位时间内数据中的新 ID 实际上只占总数据中很少的一部分，HugeCTR 发表论文中也提到相似的结论：只有一小部分的 ID 会被频繁访问。由此，我们想到先正常使用 GPU embedding 进行训练，在显存放满时再将 ID 批量驱逐至 PS 。

文章插图
根据这样的一个思路，假如 GPU embedding 里面只能存下 n 个 ID，而总 ID 有 N 个，甚至无穷多个。可以将全局的 ID 按顺序映射到 0、1、2、3…，并把把映射关系存在一个叫 ID transform 的结构中，让 GPU embedding 利用映射的结果进行正常的训练。当 GPU embedding 放满了，也就是 ID transformer 中 n 对映射的时候，再批量驱逐 ID 至 PS 。