关于开发者: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