一块GPU训练TB级推荐模型不是梦，OneEmbedding性能一骑绝尘

来源 | 机器之心

个性化推荐已成为人们获取信息的主要形式。以往，人们更多通过主动搜寻自己感兴趣的信息，而现在，基于算法推荐技术的信息分发平台会自动识别用户兴趣，快速筛选信息，推送用户所感兴趣的信息。

一方面，推荐系统大幅提升了用户体验，另一方面，个性化分发信息更精准、高效，可以帮助平台更准确地匹配用户和信息，大大提高流量变现效率，基于推荐技术的流量变现引擎甚至成就了万亿市值的庞大商业帝国。

从短视频信息流推荐、广告搜索到线上购物，这些应用都构筑于精准的推荐系统之上，背后的核心功臣就是深度学习模型。

不过，随着海量数据的积累以及更加频繁的用户数据迭代，底层系统可扩展性和训练速度面临严峻的挑战。人们发现，通用深度学习框架都不能直接满足工业级推荐系统的需求，而是必须基于通用深度学习框架做深度定制，甚至于要开发专门的系统才行。

针对现代推荐系统的种种痛点，OneFlow 团队推出了一款高性能、可扩展、灵活度高的推荐系统组件 OneEmbedding。它的使用方式和通用深度学习框架一样简单，性能却远超通用框架，甚至超过了 NVIDIA HugeCTR 这样为推荐场景定制开发的系统。

具体而言，在 DCN、DeepFM 两个模型上，无论是 FP32 还是混合精度（automatic mixed-precision, AMP）训练，OneEmbedding 的性能大幅超过 HugeCTR，而在 HugeCTR 深度优化以至于有点 “过拟合” 的 DLRM 模型上，OneEmbedding 性能与 HugeCTR 基本持平。

（以上测试环境均为：CPU Intel(R) Xeon(R) Platinum 8336C CPU @ 2.30GHz * 2；CPU Memory 1920GB；GPU NVIDIA A100-SXM-80GB * 8；SSD Intel SSD D7P5510 Series 3.84TB * 4）

当用户使用 OneFlow 搭建推荐模型时，只需使用以下数行代码对 Embedding 词表进行配置即可训练含有 TB 级别词表的推荐模型：

# self.embedding = nn.Embedding(vocab_size, embedding_vec_size)
self.embedding = flow.one_embedding.MultiTableEmbedding(
                     "sparse_embedding",
                     embedding_dim=embedding_vec_size,
                     dtype=flow.float,
                     key_type=flow.int64,
                     tables=tables,
                     store_options=store_options,
                 )

基于 OneEmbedding 搭建的常见搜索推荐广告模型案例地址：https://github.com/Oneflow-Inc/models/tree/main/RecommenderSystems

1

大规模推荐系统的挑战

一般而言，推荐系统需要使用类似性别、年龄、行为等方面的离散特征（sparse feature），在一个 Embedding 词表中用特征 ID 进行查表（lookup），取得对应的 Embedding 向量并送到下游使用。

常用的公开数据集 Criteo1T 中大概包含十亿个特征 ID，如果 embedding_dims 配置为 128，那总共需要 512 GB 空间来容纳 Embedding 参数，如果使用 Adam 优化器，由于需要保存额外的两个状态变量 m 和 v，所需存储容量就增加到 1536 GB。实际应用场景中，数据规模比 Criteo 还要高出几个数量级，模型的容量就更大了。

大规模推荐系统的核心问题就是，如何高效经济地支持大规模 Embedding 的查询和更新。权衡规模、成本和效率，出现了如下三种常见的解决方案。