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

文章插图

在这种设计下，可以使得 PS 很少介入，只有在驱逐时才需要 GPU worker 和 PS 通信。

文章插图

除此之外，这样的设计中 PS 只需要作为 KV，不需要支持参数更新，也就不需要实现优化器相关的操作，从而让 PS 团队专注于存储相关的工作。我们也支持实现了任意 KV 存储的插件，在开源版本中更是内置了 redis 插件，让 Redis 也可以作为一个 PS 来使用。

文章插图
下面介绍一些 dynamic embedding 中的设计细节。我们实现的最简基础的 ID Transformer，其实也就是用一个哈希表，使用的是 PyTorch 里高性能的 ska::flat_hash_map 。

文章插图

ID Transformer 作为流程中仅有的 CPU 操作，对性能要求可能会比较高，所以我们还实现了一个高性能的版本，以 L1 cacheline 为单位存储，从而进一步提升内存的访存效率。

文章插图

另外，对于驱逐方案，我们希望在不增加内存缓存压力的情况下，高效地融合 LRU 和 LFU 。受到 Redis 的 LFU 方案的启发，我们设计了一种概率的算法：只存储 ID 访问频数的指数。比如访问了 32 次即存储 5 。在更新频数时，如果又访问到这个 ID，就生成 5 位的随机数，如果在 5 位全为 0，也就是发生了概率为 1/ 32 的事件，我们就增加频数指数为 6 。通过这样的概率算法，就可以把 LRU 和LFU 的频数放到 uint32 里面，在不提高访存压力的情况下融合了 LRU 和 LFU 。

文章插图

最后来简单介绍一下我们的多卡方案。我们目前是将所有卡的数据都先 gather 到卡一的 ID Transformer 上，之后再 broadcast 回去。因为我们实现的 ID Transformer 的性能非常高，而且可以和 GPU 计算 Pipeline 起来，不会成为具体的性能瓶颈。

文章插图

以上就是 dynamic embedding 在设计上一些想法。在我们内部的一个万亿级的业务上，在对齐精度情况下，dynamic embedding 方案相对于我们内部原始的 GPU Tensorflow 框架有 3 倍左右的性能提升。相比于 TF 优化版也仍然有 50% 以上的性能优势。

文章插图

最后推荐大家去试用一下 Torchrec 。对于相对较小的业务，比如百亿下的业务，推荐大家直接使用原生的 TorchRec：即插即用，不需要任何的二次开发，性能可以得到成倍的提升。对于极大的业务，则推荐大家尝试配合我们合进 TorchRec 的 dynamic embedding，一方面方便连接内部的 PS，另一方面也支持 embedding 的扩展和渐进迁移，同时还是可以获得一定的性能提升。