Hybrid Embedding：蚂蚁集团万亿参数稀疏 CTR 模型解决方案

文章介绍了蚂蚁AIInfra团队在TFPlus项目中提出的HybridEmbedding方案，以应对大规模稀疏CTR模型的存储和通信开销。该方案结合NVMeSSD和层次化存储，解决了传统内存参数服务器的不足，通过冷热特征划分和动态缓存管理优化了资源使用。HybridEmbedding已在蚂蚁推荐场景中应用，显著节省了内存资源。

导读

稀疏 CTR 模型是用于预测点击率的一类深度学习模型的统称。通过处理高维稀疏特征数据实现高效的广告点击率预测，准确的 CTR 预测可以更好地进行广告投放决策，提升广告投放效果。在稀疏 CTR 模型场景下，通常使用大量高维稀疏特征，这类 Embedding 矩阵造成了巨大的存储开销。本文介绍了蚂蚁 AI Infra 团队在 TFPlus 项目中提出的 Hybrid Embedding 方案，以应对海量 Embedding 带来的存储资源、特征治理方面的挑战。

背景

在实际场景中，稀疏 CTR 模型会使用大量高维稀疏特征，如商品 ID、用户 ID 等，这些特征的取值可能非常稀疏。处理这样的稀疏 Embedding 矩阵需要引入巨量的存储和通信开销。实际场景中，稀疏层参数可达千亿甚至万亿，给系统设计带来巨大挑战。

由于传统的 TensorFlow Variable 存在静态 Shape 难以预估、特征冲突等问题，TFPlus 引擎基于 TensorFlow 的自定义算子的扩展能力，实现了 KvVariable 内存参数服务器。KvVariable 采用分布式哈希存储方案，针对稀疏 CTR 模型的超大 Embedding 层进行了一系列优化，如配合 DLRover 支持 Embedding 层的弹性扩缩容，增量更新方式提升模型上线速度，提供了特征裁剪、特征淘汰等能力。

在蚂蚁的推荐场景，KvVariable 内存参数服务器得到了广泛的应用，具体方案如下：

训练：采用了数据并行和模型并行的方式，并基于“有分片，无副本” （With-sharding，No-replica） 架构。多个 worker 节点进行异步的数据并行计算，负责执行深度神经网络的前向传播和反向传播计算；多个 ps 节点进行模型并行计算，负责稀疏 Embedding 参数的分片存储、收集汇总梯度以及参数的更新。
推理：采用了“无分片、有副本” （No-sharding，With-replicas） 架构。在这种架构中，不对稀疏 CTR 模型的