VLDB 2023 | 清华大学快手联合提出首个持久性内存参数服务器系统

论文地址：https://www.vldb.org/pvldb/vol16/p1013-xie.pdf

作者：谢旻晖（清华大学）、陆游游*（清华大学）、汪庆（清华大学）、冯杨洋（清华大学）、柳嘉强（快手）、任恺（快手）、舒继武（清华大学）

与计算机视觉、自然语言处理等传统深度学习场景不同，推荐、搜索、广告等业务的输入包含大量的 ID 类型特征（如用户 ID、视频 ID）。这些 ID 类型特征具有极高的维度与很强的稀疏性，无法被深度神经网络直接学习。如图1左部所示，为了解决这个问题，嵌入模型利用多个嵌入表（embedding table）将 ID 类特征转化为神经网络可学习的低维稠密向量（称为嵌入向量，embedding vectors，简称EMB)，之后经过池化操作后作为 DNN 模型的输入，并最终参与到对模型目标的预估中。

近年来，随着数据规模与特征数量的迅速增长，工业级嵌入模型中的参数量已达万亿级。其中来自于嵌入表中所含的参数占绝大部份，DNN 部分占比较少。面对如此庞大的规模，现有预估系统往往采用存算分离架构（如图 1右部所示）。具体而言，由一组大内存 (DRAM) 的参数服务器 (Parameter Server，PS) 分片存储嵌入表的参数，并为推理服务（Inference Server）提供实时的 EMB 参数查询。PS从预估服务接收特征ID，查找索引(Indexing) 得到对应 EMB 参数的地址，再将它们聚集（Gathering）并序列化后返回。

图 1 嵌入模型结构与典型预估系统架构

然而，随着模型参数量的不断膨胀，传统DRAM参数服务器带来的存储成本高、宕机恢复时间长等问题日益严重。新型持久性内存 (Persistent Memory, PM) 以其大容量、低成本和持久性的优点为参数服务器系统设计提供了新的机遇。但利用PM构建高效的参数服务器仍面临着两方面挑战：

（1）如何减少PM高读延迟对模型推理性能的影响？PM的读延迟是DRAM的3倍，如果在系统层处理不当，其高延迟可能被完全暴露给应用层，导致模型推理性能下降，并进一步影响服务SLA 和用户体验。

（2）如何处理CPU 资源瓶颈对参数服务器性能的限制？单台PM参数服务器上可存储的参数量比DRAM大8倍，这意味着系统必须在更慢的介质上服务更多参数访问请求，但只使用相同的CPU资源。这使得CPU成为基于PM的参数服务器的性能瓶颈。

（1）访问读密集。PS中的大部分操作（约98%）都是只读的，只有少量的写入操作（约2%）用于保证模型新鲜度。

（2）容量负载稳定。PS的参数容量和内存使用是稳定的。尽管模型在不断接收新ID，但PS会淘汰旧嵌入参数以容纳新嵌入参数。因此，总负载始终保持稳定。

（3）单一请求批处理。一个模型推理请求通常涉及数百到数千个ID。

图 2研究动机实验

同时，为了深入分析参数服务器系统在 PM 上的基本性能，我们分别使用现有的 PM 哈希索引(包括 Level、Clevel、CCEH、Dash) 建立了一个 PS 系统。如图 2 (d-e)，我们发现：

（1）CPU 是主要的吞吐瓶颈。在峰值吞吐时，系统充分利用了所有的 CPU 核心（利用率为100％），但相对地，只占用了46％的网卡带宽。同时，考虑到 PM 的随机读带宽（36GB/s）远高于网卡带宽（25GB/s），我们总结出 CPU 是参数服务器首要的吞吐瓶颈。

（2）索引查找与聚集两步都耗费了大量 CPU 时间。对于不同的 PM 哈希索引，至多有71% 的 CPU 时间被消耗在索引查找上。除索引查找外，参数服务器的 CPU 在聚集步骤上也耗费了多达29%的时间。我们将其归因于单一请求的批次过大与 PM 的高读取延迟。

针对上述问题，我们提出了 PetPS，一个基于 PM 的参数服务器系统。PetPS 并没有改变传统 PS 的“索引+聚集”的两步工作流程，它的主要创新在于：根据我们对真实模型服务的负载分析，对这两个步骤进行了适合于PM的设计和实现，从而有效解决上述挑战。

图 3 PetPS架构示意图

如图 3 所示，PetPS 包含两个核心组件：

（1）针对索引部分，我们提出 PetHash 索引，一个为 PS 场景高度优化的持久性内存哈希表，其核心设计思想是结合嵌入模型独特的负载特征，最大限度地减少 PM 的读访问次数。

（2）针对聚集部分，我们提出网卡聚集机制 （NIC Gathering），将 CPU 的聚集操作卸载到网卡上，利用网卡的 DMA 引擎聚集参数，以规避PM高读取延迟对 CPU 流水线的影响。

为了减少 PM 高读延迟的影响，我们根据 PS 的容量负载稳定、数据存在热点、单一请求批处理三个特点，定制了一个 PM 上的哈希索引 PetHash，其包含单层结构、热点感知的迁移机制、预取机制三点优化，以尽可能最少化参数访问路径上的PM读取次数。

（1）单层结构。现有PM哈希索引通常采用多层级结构，以减少 rehash 的开销，但这种方式导致单次索引查找涉及多次PM读操作。例如，目前最先进的PM索引 Dash ，需要多次指针追逐操作来定位一个键值对（目录→段→桶→键值对）。由于我们的 PS 具有容量负载稳定性，不需要 rehash 操作，PetHash 设计了一个单层结构来降低读成本，其好处是单次索引查找最优情况下只有一次 PM读操作。

具体而言，如图 4所示，PetHash 由一些桶（Bucket）以单层的方式组织形成，使用开放地址法解决哈希冲突。每个桶大小为256字节（设为PM硬件内部块大小），内部包含FP (指纹机制)、version（用于实现无锁查询）、overflow (记录本应被放入当前桶，但由于当前桶已满而被移入其他桶的键值对数量)等元数据信息。对于读操作，PetHash会从主桶开始逐个探测，直到找到该KV，或者遇到一个overflow为0的桶，此时意味着当前KV不存在于哈希表中。PetHash的插入操作、删除操作、元数据修改、并发控制等其他细节可以阅读我们的论文。

图 4 PetHash 结构示意图

（2）热点感知的迁移机制。我们定义每个 KV 的主桶（home bucket）为索引查找路径上的首个桶，由 Key 的哈希值决定。真实世界嵌入模型中的访问往往存在热点，热点感知的迁移机制优先将热键值对放置在对应的主桶中，这样查找某个桶仅需一次 PM 读操作（符合 PetHash 的核心设计原则）。具体而言，PetHash 使用一个专门的迁移线程来移动热键值对。迁移线程定期检查每个热点 KV 是否在自己的主桶中，若不在，迁移线程首先会将它插入主桶中，如若主桶已满，桶内一个较冷的键值对可能被移出以腾出空间（递归地执行迁移操作），最后将该热键值对从原位置中删除。

（3）预取机制。一个 PS 请求通常包含数百个参数的读写操作，且这些均由参数服务器的一个线程串行查询。这种批查询的方式为 PetHash 提供了预取的机会。在对当前 KV 进行索引之前，PetHash 会发出预取指令以预取下一个 Key 对应主桶的数据至 CPU 缓存。因此，叠加热点感知的迁移机制，下一个键的索引过程有极大概率完全命中 CPU 缓存，从而隐藏大部分的 PM 读延迟。

完成索引操作后，PS需要根据查询结果聚集参数，并返回给客户端。

传统基于CPU的聚集方式（CPU Gathering）会占用多达 30% 的 CPU 资源。原因是在 CPU 流水线中引入了大量高延迟的 PM 随机读，导致严重的 CPU 缓存缺失，进而浪费 CPU 指令周期。

我们提出网卡聚集机制（NIC Gathering），其核心观察是目前商用数据中心网卡通常提供了 Scatter-Gather DMA（SG-DMA）功能，可用于聚集PM上的参数。图 5 对比了传统基于 CPU 的聚集方式与 NIC Gathering 机制。在 NIC Gathering 机制中，CPU 数据路径上无需接触 PM 上的参数本身，仅需向网卡发送 SG-DMA 请求，这避免了 CPU 访问低速 PM 介质时的停顿，从而提升 PS 的吞吐。

图 5 传统基于 CPU 的聚集方式与 PetPS 网卡聚集（NIC Gathering）机制对比

由于 DMA 是由网卡异步完成的，PetPS 必须确保在 DMA 过程中的两个不变式：(i) 正在被 DMA 的嵌入参数不得被修改；(ii) 其对应内存空间不得被释放。否则，被 DMA 聚集的参数可能出现版本不一致的情况。为此，PetPS 使用了写时拷贝（copy-on-write）与我们提出的 epoch list-based reclamation 机制分别保护这两个不变式。具体细节可以阅读我们的论文。

（1）PSLite，来自 MXNet 的经典参数服务器系统实现。

（2）DashPS，使用目前 SOTA 持久性内存哈希索引 Dash 替换了PSLite的索引。

（3）KuaiPS，来自快手内部对 DRAM 参数服务器的实现。

其他详细的实验设置可以参考论文。

A. 端到端性能

图 6 左图：端到端吞吐量；右图：吞吐-延迟曲线

图 6 左展示了每个系统的峰值吞吐。单位 op 指的是一次参数服务器的 pull 或 push 操作。与其他系统相比，PetPS 在 YCSB-C 负载上的峰值吞吐提升 1.3-1.6 倍，在 Production 负载上提升可达 1.4-1.7 倍。

图 6 右展示了系统的吞吐-延迟曲线。在低请求压力下（如 20Mops/s），PetPS 的中位数和 P99 延迟分别降低了 2.9-5.5 倍和 3.1-5.1 倍。在峰值吞吐下，中位数和 P99 延迟降低了 2.6-3.1 倍和 2.2-3.3 倍。上述数据整体说明了 PetPS 设计的有效性。

B. PetHash性能

图 7 左图：索引吞吐量；右图：各索引平均每个操作访问 PM 的字节数

我们使用 YCSB-C 负载来评估 PetHash 和现有 PM 哈希索引的性能。我们评估了三种分布，包括均匀分布、Zipfian-0.9 和 Zipfian-0.99（数值越大意味着分布越倾斜）。由于 PetHash 的设计减少了PM读操作次数（至多达 2 倍，见图 7右），相比于现有的 PM 哈希索引，PetHash 的吞吐量增加了 1.3-2.5 倍。与原始 PetHash 相比，使用热点感知的迁移机制缩短了热点 KV 的查询路径，从而提升了 1.1 倍吞吐，且该机制在更倾斜的分布下带来了更多提升。Prefetch 通过隐藏 PM 读取延迟，可以进一步将吞吐提升11-40%。

C. NIC Gathering性能

图 8 左图：两种聚集方式下CPU time breakdown；右图：两种聚集方式下端到端吞吐与延迟

图 8 左展示了峰值吞吐量下，PetPS 分别使用两种聚集方式时索引查找和参数聚集的耗时。图 8 右展示了在两种聚集方式下的端到端吞吐、中位数延迟、P99 延迟。我们发现 NIC Gathering 机制可以将 CPU 的聚集时间从 180 微秒降低至 14 微秒，省下的 CPU 开销使得 PS 端到端吞吐提升多达 1.2 倍，同时中位数与 P99 延迟也分别降低 2.3 倍与 2 倍。

我们也在 DRAM 上测试了 NIC Gathering 的性能。有趣的是，在 DRAM 上使用 NIC Gathering 反而会导致 30% 吞吐量的下降。为细究其原因，我们做了如下的实验。如图 9所示，我们设置两次 Gathering 操作之间的间隔时间（x轴），对比两种聚集方式的吞吐量（y轴）。随着发起 Gathering 的频率越来越高，NIC Gathering 性能很快进入瓶颈，而 CPU Gathering 反而可以达到更高的吞吐。对于PS来说，索引部分的时间就是间隔时间，我们统计了在 DRAM 与 PM 上的索引时间，发现它们恰好位于交叉点的两侧。这就是 NIC Gathering 为何在 PM/DRAM 上表现不同的原因。

图 9 ：在 DRAM 上 NIC Gathering 性能测试

PetPS 已部署于快手推荐业务，为行业内首次在推荐系统中应用实践 PM 技术，成功应对了每日超百亿次视频推荐的访问压力。根据我们内部的成本分析，与 DRAM 参数服务器相比，PetPS 可以在不影响服务性能的同时，降低约 30% 的机器成本 (TCO)。

本工作由清华大学计算机系存储实验室与快手合作完成，是继Kraken [2]、 Fleche [3]之后在Storage for AI方向的进一步探索。PetPS目前已开源于GitHub (https://github.com/thustorage/PetPS)。

本文撰稿作者：清华大学谢旻晖

参考文献

[1] Xie, Minhui, Youyou Lu, Qing Wang, Yangyang Feng, Jiaqiang Liu, Kai Ren, and Jiwu Shu. "PetPS: Supporting Huge Embedding Models with Persistent Memory." Proceedings of the VLDB Endowment 16, no. 5 (2023): 1013-1022.

[2] Xie, Minhui, Kai Ren, Youyou Lu, Guangxu Yang, Qingxing Xu, Bihai Wu, Jiazhen Lin, Hongbo Ao, Wanhong Xu, and Jiwu Shu. "Kraken: memory-efficient continual learning for large-scale real-time recommendations." In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-17. IEEE, 2020.

[3] Xie, Minhui, Youyou Lu, Jiazhen Lin, Qing Wang, Jian Gao, Kai Ren, and Jiwu Shu. "Fleche: an efficient GPU embedding cache for personalized recommendations." In Proceedings of the Seventeenth European Conference on Computer Systems, pp. 402-416. 2022.