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

导读

稀疏 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 内存参数服务器得到了广泛的应用，具体方案如下：

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

2. 推理：采用了“无分片、有副本” （No-sharding，With-replicas） 架构。在这种架构中，不对稀疏 CTR 模型的 Embedding 进行分片存储，而是使用多个副本 （replicas） 来存储相同的数据。每个副本都包含完整的 Embedding 数据，并且可以独立地进行推理，从而提高系统的可靠性。

KvVariable 内存参数服务器方案通过分布式的扩展能力，一定程度上解决了单机处理海量特征的瓶颈问题，但没有很好地解决稀疏 CTR 模型对于内存资源日益增长的需求。同时，采用参数服务器的方式增加集群规模时，会引入额外的通信开销和稳定性下降，下面将具体分析。

问题分析

分布式内存参数服务器一定程度上解决了单机处理海量特征的瓶颈问题，但仍然存在很多不足：

• 通信：增加存储节点，意味着系统通信成本的增加 （线性关系）
• 稳定性：增加存储节点，意味着系统稳定性的降低 （指数关系）

此外，在蚂蚁的推荐场景中，受集群资源总量限制和业务预算限制，资源问题是一个比较大的痛点：

• 有限的资源：资源总量受限于 ROI 和业务预算。这意味着系统可能无法支持超出资源限制的模型的训练和推理部署。更大的模型、更多的参数可以实现对商品和用户更精细的刻画，从而获得更好的推荐效果。但是，由于资源的限制，推荐系统可能无法支持更大模型的训练和推理。

• 成本挑战：稀疏 CTR 模型需要海量存储，内存方案比较昂贵，希望能够寻求更具性价比的存储方案实现降本增效。

• 资源难以预估：在线学习场景中，由于数据分布的变化，模型可能会持续增长，导致系统内存占用不断增加，进而导致资源用量超出预先设定的配额，出现内存不足 (OOM) 的问题。这将对推荐系统的稳定性和效果产生负面影响。

• 特征治理：Embedding 作为一种特征数据可能被多任务使用，如果每个任务都维护一份完整的 Embedding 副本，将导致数据冗余和资源浪费的问题。

Hybrid Embedding

全新设计的 Embedding 存储架构

基于以上问题，AI Infra 团队提出了一种基于多层存储的 Hybrid Embedding 方案，对 KvVariable 参数服务器进行了极致优化，以满足日益增长的稀疏 Embedding 的存储需求，并提高训练稀疏 CTR 模型的性价比。此外，通过 Embedding Service，解决了 Embedding 特征的冗余问题，满足了特征治理的需求。下面将重点介绍 Hybrid Embedding 中的关键技术。

NVMe SSD

NVMe 等高速读写设备的发展对解决推荐系统内存资源问题提供了有力的支持，相较于传统的 SATA 存储设备，NVMe 设备具有 10 倍的读写速度，随机读写可达每秒 900K/275K IOPS，PCIe Gen4 下理论带宽 7.8GB/s，延迟 80us。

Embedding Service

Embedding Service 是提供一站式高扩展性嵌入式服务的技术，通常用于将离散的输入特征 （如文本、用户 ID、商品 ID 等） 映射到对应的向量。它作为一个独立的服务嵌入到机器学习模型中，供模型在线训练和推理使用。用户事先通过离线训练或在线学习的方式生成表征向量，存储在 Embedding Service 中。同一个 Embedding 向量可被多任务或多副本共享使用，避免数据冗余和资源浪费。

层次化存储与冷热特征

层次化存储思想基于数据的访问频率和访问速度的不同，将数据划分为多个层级，并根据不同层级的存储介质和访问速度，将数据存储在不同的层级中。通常，存储介质的容量与访存速度成反比，较高层级的存储介质 （如内存、高速缓存等） 访问速度较快但容量较小，适合存储频繁访问的热数据；而较低层级的存储介质 （如磁盘、分布式文件系统等） 容量较大但访问速度较慢，适合存储访问较少的冷数据。

稀疏 Embedding 特征在训练中通常服从幂律分布，低频特征占比大，高频特征占比小。因此，可以根据频次信息对 Embedding 进行冷热分组，分配在不同的存储介质中。

技术选型

为了保证层次化存储的读写性能，高效地利用 NVMe SSD 实现随机读写很重要。为此，我们对比了业界多款 SSD KV 存储，针对广告推荐模型模型训练到一定程度后特征分布趋于稳定，存在大量随机读、更新操作的特点进行分析。

1. 读方面：LevelDB、RocksDB 等采用常规 LSM Tree+SSTable 结构的存储，随机读需要进行层次访问文件与 O(log N) 的查询，性能存在一定瓶颈；采用 B+tree 的存储由于层级较少，虽然需要一定次数的查询，对比 LSM Tree 结构查询性能更好且稳定。

2. 更新：常规 LSM Tree 结构的存储定期写入磁盘，有较好的写入性能；而 B+tree 在写入修改时需要大量移位操作。

最终我们选择了一款蚂蚁内部自研的高性能 SSD 存储引擎 PHStore，其能够提供高效的 SSD 随机读写能力，在保证写入性能的同时，随机读性能远优于上述存储架构。在提供高读写性能的同时，该存储引擎还有高可用、高可靠的容错机制以及线性扩展能力，使得我们能够方便地基于其提供的分布式服务构建 Embedding Service。

重点工作

1. 基于 LFU 算法依据频次信息实现动态的冷热特征划分，并依据时段内特征查询的 IOPS 与缓存命中率计算缓存容量，从而实现动态的缓存更新策略。

2. 我们设计了一套索引 Meta 表的结构，以实现特征的精准查询，采用 Lazy Load 的方式更新缓存，定期更新时段内发生变化的增量 IDs，尽量平衡缓存命中率与冷热表置换对训练性能的影响。

3. 对 Embedding 查询进行了一系列优化，在前反向传播过程中对参数进行拼接、缓存，实现一次传播一次查询，成倍减少 SSD 的访问频次。

4. 在模型容错方面，实现了混合存储的动态分片，均匀地将参数量与 IO 负载分布在存储节点上，当即将发生 OOM 或到达 IO 瓶颈时，配合 DLRover 的调度进行无感知的扩容。

效果与总结

我们对比了业内常见的模型结构，如 DeepFM、DLRM、Transformer 等，均在性能不明显下降的情况下达到明显的内存优化。

在蚂蚁推荐场景中，目前，Hybrid Embedding 已全面应用于在线学习场景，在性能几乎与 DRAM kv 存储打平的同时，大部分原本被 Embedding 占用的内存资源被优化，ps 节点内存平均节省 50%左右。

本文介绍了蚂蚁推荐系统 DLRover 采用的大规模稀疏存储 Hybrid Embedding 方案，该方案不仅实现了在系统应用中降低成本和提高效率的目标，还为不断增长的模型参数规模打下了技术基础。未来，我们将继续优化 Hybrid Embedding 以满足高性能训练需求，将其应用于更多的业务场景中。

TFPlus 项目后续会协同大模型技术栈的训练加速方案一起在 DLRover 项目中开源，敬请关注：https://github.com/intelligent-machine-learning/dlrover