kuaishou在《TWIN: TWo-stage Interest Network for Lifelong User Behavior Modeling in CTR Prediction at Kuaishou》中提出了TWIN的长序列建模方法。

摘要

终身用户行为建模（Life-long user behavior modeling），即从数月甚至数年的丰富历史行为中提取用户的隐藏兴趣，在现代CTR预测系统中起着核心作用。传统算法大多遵循两个级联阶段：一个简单的通用搜索单元（GSU），用于快速和粗略地搜索数万个长期行为，以及一个精确搜索单元（ESU），用于对GSU的少数最终选手进行有效的目标关注（TA）。尽管高效，现有算法大多存在一个关键限制：GSU和ESU之间的目标-行为相关度度量不一致。因此，它们的GSU通常会错过高度相关的行为，但会检索被ESU认为不相关的行为。在这种情况下，ESU中的TA，无论如何分配注意力，都会偏离真实的用户兴趣，从而降低整体CTR预测精度。为解决这种不一致性，我们提出了TWo-stage Interest Network（TWIN），其中我们的Consistency-Preserved GSU（CP-GSU）采用与ESU中TA相同的目标-行为相关度量，使两个阶段成为孪生。具体而言，为了打破TA的计算瓶颈并将其从ESU扩展到GSU，即从行为长度102扩展到长度104-105，我们通过行为特征分割构建了一种新的注意机制。

• 对于行为的视频固有特征，我们通过高效的预计算和缓存策略计算它们的线性投影。
• 对于用户-物品交叉特征，我们将每个特征压缩为注意力分数计算中的一维偏置项，以节省计算成本。

两个阶段之间的一致性，加上CP-GSU中有效的TA-based相关度量，为CTR预测的显著性能提升做出了贡献。在快手的460亿规模的真实生产数据集上进行的离线实验和在线A / B测试表明，TWIN优于所有比较的SOTA算法。通过优化在线基础设施，我们将计算瓶颈降低了99.3％，这有助于TWIN在快手上的成功部署，每天为数亿活跃用户提供主要流量服务。

1.介绍

作为中国最受欢迎的短视频分享应用之一，快手强烈依赖于其强大的推荐系统（RS）。每天，RS帮助数亿活跃用户过滤掉数百万个不感兴趣的视频，找到他们感兴趣的内容，留下数十亿的点击日志。这些巨大的数据不仅为RS的训练提供了数据支持，而且推动了技术革命，不断提升了这个平台的用户体验和业务效果。

在现代RS中，一个基本任务是点击率（CTR）预测，旨在预测用户点击一个item/视频的概率[2,10,32]。准确的CTR预测可以指导RS为每个用户提供其喜欢的内容，并将每个视频传递给其感兴趣的受众。为了实现这一目标，CTR模型应该高度个性化，并充分利用稀缺的用户信息。因此，终身用户行为建模，即从丰富的长期历史行为中提取用户的隐藏兴趣，通常作为CTR模型的关键组成部分[7,16,34-36]。

工业终身行为建模算法大多遵循两个级联阶段[19]：（1）通用搜索单元（GSU），对数万个长期行为进行快速粗略搜索，并输出最相关的少数目标行为；（2）精确搜索单元（ESU），对来自GSU的最终候选进行有效的目标关注（TA：Target Attention）。这种两阶段设计的原因有两个原因：

• 一方面，为了准确捕捉用户的兴趣，TA是强调目标相关行为和抑制目标不相关行为的合适选择
• 另一方面，TA的高昂计算成本限制了其适用的序列长度最多只有几百个。为此，一个简单快速的GSU作为预过滤器对于截断在短短几个月内就可以轻松达到$10^4-10^5$的工业规模行为序列至关重要。

近年来，出现了许多关于两阶段终身行为建模的新兴研究，它们的主要区别在于GSU策略，即如何粗略选择目标相关行为。例如：

• SIM Hard [19]：仅从与target item相同的类别中选择行为
• SIM Soft [19]：通过内积计算预训练item embedding的目标-行为相关度分数，并选择相关度最高的行为
• ETA：使用局部敏感哈希（LSH）和汉明距离来近似计算相关度分数[3]
• SDIM：通过多轮哈希碰撞从具有相同哈希签名的行为中采样目标行为，等等。

尽管已经广泛研究，现有的两阶段终身行为建模算法仍然存在一个关键限制：GSU和ESU之间的不一致性（如图11所示）。具体而言，GSU使用的目标-行为相关度量既粗略又与ESU中的TA不一致。因此，GSU可能会错过相关的行为，但会检索被ESU认为不相关的行为，浪费ESU宝贵的计算资源。在这种情况下，ESU中的TA，无论如何分配注意力，都会偏离真实的用户兴趣，从而降低整体CTR预测精度。

为了解决这种不一致性，我们提出了TWIN：TWo-stage Interest Network，用于终身用户行为建模，其中Consistency-Preserved GSU（CP-GSU）采用与ESU中TA相同的目标-行为相关度量，使两个阶段成为孪生。为了将昂贵的TA扩展到CP-GSU中，TWIN通过有效的行为特征分割、简化的TA架构和高度优化的在线基础设施打破了TA的关键计算瓶颈，即所有行为的线性投影。具体而言，对于行为的视频固有特征（例如视频ID、作者、持续时间、主题），这些特征在用户/行为序列之间共享，我们通过高效的预计算和缓存策略加速它们的投影。对于行为的用户-视频交叉特征（例如用户的点击时间戳、播放时间、评分），其中缓存不适用，我们通过将它们的投影压缩为偏置项来简化TA架构。通过优化在线基础设施，我们成功将TA的适用序列长度从ESU中的$10^2$扩展到CP-GSU中的$10^4 ~ 10^5$。两个阶段之间的一致性，加上CP-GSU中有效的基于TA的相关度量，为CTR预测的显著性能提升做出了贡献。

主要贡献：

• 在我们提出的TWIN中，CP-GSU精确而一致地检索不仅与目标相关，而且ESU认为重要的行为，最大化行为建模的检索效果。据我们所知，我们是第一个成功解决两阶段终身行为建模问题不一致性的团队。
• 我们通过在快手的460亿规模的工业数据集上进行大量离线实验和在线A/B测试来验证TWIN的有效性。我们通过消融研究验证了我们的有效性，并展示了TWIN带来的显著在线收益。
• 我们构建了高效的工业基础设施，将TWIN应用于实际在线RS。我们提出了有效的预计算和缓存策略，将TWIN的计算瓶颈，即CP-GSU中行为的线性投影，降低了99.3％，并满足了在线服务系统的低延迟要求。TWIN现已部署在快手的RS上，每天为3.46亿活跃用户的主要流量提供服务。

2.相关工作

我们的工作与两个活跃的研究领域密切相关：CTR预测和长期用户行为建模。

2.1 点击率预测

CTR预测旨在预测用户的个性化兴趣，对于现代RS至关重要。早期的CTR模型是浅层的，主要关注于利用特征交互，例如因子分解机（FM）[22]和场感知因子分解机（FFM）[12]。随着深度学习的成功，深度CTR模型得到广泛研究并成为主流选择。例如，陈等人[2]和张等人[33]首次将深度模型应用于CTR任务。Wide&Deep [5]结合了宽线性模型和深度模型，充分利用特征交互的记忆和深度架构的泛化优势。DeepFM [10]和DCN [26,27]改进了Wide&Deep的宽部分，以增加特征交互能力。xDeepFM [15]和AFM [29]进一步利用类卷积层和注意机制来改进深度部分并提高模型性能。

随着CTR模型变得越来越个性化，用户行为建模，即从历史行为的总结中捕捉用户的隐藏兴趣，成为一个关键模块。由于计算资源的限制，早期的算法大多采用目标无关的方式，因此可以在离线情况下高效地预计算[8,23,31]。为了更好地提取用户对特定item的兴趣，采用了各种TA机制。DIN [36]通过历史行为上的TA表示用户兴趣，强调目标相关行为。DIEN [35]进一步使用ARGRU（经典GRU的基于注意力的变体）引入行为之间的时间关系。DSIN [9]将行为分为多个会话，并在每个会话内进行自注意力计算，以强调会话内关系。MIND [14]和DMIN [30]通过多个向量表示用户兴趣。BST [4]、SASRec [13]和BERT4Rec [24]也使用变压器来提高模型的性能和并行性。

2.2 Long-Term User Behavior Modeling

随着TA和兴趣建模在现代工业RS中的有效性得到确认，研究人员开始对越来越长的行为进行建模。Liu和Zamanian [16]将长期和短期兴趣结合在CTR预测中。MIMN [18]将用户行为存储为用户兴趣中心（UIC）的记忆矩阵，并在新的用户行为到来时更新记忆。然而，MIMN难以扩展到长度超过$10^3$的序列，并为不同的候选项生成相同的记忆矩阵，携带无用的噪声并损害TA。

最近，SIM [19]和UBR4CTR [20,21]引入了两阶段级联框架来解决这些挑战，并在CTR预测中实现了SOTA性能。传统的两阶段算法通常由以下两部分组成：

• 1）一个简单快速的GSU，从数千个用户行为中检索与目标项最“相关”的item
• 2）一个注意力ESU，对GSU的最终候选item执行TA

UBR4CTR在其第一阶段中使用BM25作为相关度量。而在原始的SIM中，有两个具有不同GSU设计的实例。SIM Hard的GSU从与目标项相同的类别中选择相关项，而SIM Soft的GSU使用预训练item embedding的内积作为相关度量。尽管两阶段设计迈出了重要一步，但原始的GSU仍面临着高计算负担，并且与ESU具有不同的检索度量，导致两个阶段之间的不一致性。

最近，ETA [3]使用局部敏感哈希（LSH）对由ESU训练的item embedding进行编码，并通过汉明距离（HD）从长期行为中检索相关项。SDIM [1]通过多轮哈希碰撞从具有相同哈希签名的行为项中采样target item，并通过线性聚合这些采样的行为项来获取用户兴趣。ETA和SDIM采用End2End训练是积极的。换句话说，它们的两个阶段共享相同的embedding。然而，在检索策略方面仍存在不一致性，具体而言是网络结构和参数。

在本文中，我们提出将TA结构扩展到GSU，并将embedding和attention参数从ESU同步到GSU，保持端到端训练。结果，在网络结构和模型参数方面实现了一致性，相比于ETA和SDIM，获得了显著的性能提升。我们在表1中详细说明了我们的模型与其他模型的差异。请注意，我们的工作与旨在加速变压器（例如LISA [28]）的索引算法不同。它们通过将行为映射到码本并查找距离来近似相关度量计算。而我们的工作以及许多其他两阶段算法使用精确的距离计算，但使用GSU作为预过滤器来减少行为数量。

3 TWIN在快手CTR预测中的应用

首先，在第3.1节中，我们回顾了CTR预测问题的一般基础知识。然后，在第3.2节中，我们描述了快手CTR预测系统的模型架构。接着，在第3.3节中，我们进一步深入探讨了我们提出的保持一致性的终身用户行为建模模块——两阶段兴趣网络（TWIN）。最后，在第3.4节中，我们介绍了必要的加速策略，以确保TWIN成功部署在快手的主流量上。

所使用的符号总结在表2中。

3.1 基础知识

CTR预测的目的是：在给定特定上下文的情况下预测用户点击一个item的概率。准确的CTR预测不仅通过提供首选内容提升用户体验，而且通过吸引感兴趣的受众，有益于内容生产者和平台的业务效益。因此，CTR预测已成为各种工业RS的核心组成部分，特别是像快手这样的短视频推荐平台。

CTR预测通常被公式化为一个二元分类问题，目标是学习一个预测函数 $𝑓: R_d \rightarrow R$，给定：

• $D=\lbrace (x_1,𝑦_1), \cdots,(x_{\mid D\mid}, 𝑦_{\mid D \mid})\rbrace$: 一个训练数据集。
• $x_i \in R_d$：是第i个训练样本的特征向量（即用户、item和上下文特征的串联）
• $𝑦_i \in \lbrace 0,1 \rbrace$：是表示用户是否点击（1）该项或未点击（0）的label

预测的CTR计算公式如下：

\[𝑦^i =\sigma(𝑓(x_𝑖))\]

…(1)

其中：

• $𝜎(\cdot)$是将𝑓的预测缩放到(0,1)的sigmoid函数

模型的训练通过最小化负对数似然来完成：

\(l(D)=-\frac{1}{|D|} \sum_{𝑖=1}^{|D|} 𝑦_i log(\hat{𝑦}_i)+(1−𝑦_i)log(1−\hat{𝑦}_i)\) … (2)

为简洁起见，当不会引起混淆时，在以下各节中省略训练样本索引𝑖。

3.2 CTR预测的架构

我们现在介绍快手CTR预测系统的架构，详细信息如图2所示。

3.2.1 embedding layer

在底部，我们的模型从一个feature embedding layer开始，它会将训练样本的原始特征转换为embedding向量。

不失一般性，我们假设所有特征在必要的预处理后都是类别型。对于具有词汇表大小为$𝑣_𝐴$的特征𝐴，我们首先将分类信息编码为一个one-hot/multi-hot编码$xA,hot \in {0,1}^{𝑣_𝐴}$。例如：

\[WeekDay=Mon => x_{WeekDay,hot} = [1, 0, 0, 0, 0, 0, 0]^T， \\ Topic={Funny, Pet} => x_{Topic, hot} = [\cdots, 0, 1, 0, \cdots, 0, 1, 0...]^T\]

请注意，在大多数工业系统中，词汇表大小（特别是用户/作者/视频ID的大小）可以轻松扩展到数亿。因此，一种常见的策略是将极高维度的one-hot编码转换为低维度的嵌入向量，

\(x_{A,emb} = 𝐸_𝐴 x_{A,hot}\) …(3)

其中：

• $𝐸_𝐴 \in R^{𝑑𝐴 \times 𝑣_𝐴}$是特征𝐴的embedding字典
• $𝑑_𝐴$是embedding维度

在我们的系统中:

• 对于具有大词汇表的id特征，我们将embedding维度设置为64，
• 对于其他特征，如视频主题、视频播放时间戳，我们将embedding维度设置为8。

在所有上层中，我们将embedding向量作为输入，因此为简洁起见省略了“emb”下标。

3.2.2 深度网络

我们的CTR预测的总体架构如图2所示。

图2 快手CTR预测系统中的TWIN。与传统的两阶段行为建模算法不同，TWIN在CP-GSU和ESU中采用相同的目标-行为相关度度量，包括相同的网络架构（如左图所示）和相同的参数值（如中下部分所示）。这是具有挑战性的，因为MHTA的计算成本很高，因此只适用于ESU（具有100个行为），而不适用于CP-GSU（具有104个行为）。 我们通过提出以下方法来解决这个挑战：1）高效的特征拆分和投影策略，以不同的方式处理项的固有特征和用户-项交叉特征（如右下图所示）；2）简化的目标注意力架构，通过将交叉特征压缩为偏置项来加速目标注意力的效率（如左图所示）。

上层模块由堆叠的神经网络和ReLU组成，作为一个混合器，学习三个中间模块的输出之间的交互作用：

• TWIN，提出的保持一致性的终身用户行为建模模块，通过两个级联的行为建模子模块提取用户兴趣：
  • 1）保持一致性的一般搜索单元（CP-GSU）：从成千上万的长期历史行为中进行粗略搜索，找到100个最相关的行为；
  • 2）精确搜索单元（ESU）：对CP-GSU的100个最终选手采用attention机制，捕捉精确的用户兴趣。与通常由“轻量级”GSU和“重量级”ESU组成的传统算法不同，我们提出的CP-GSU采用与ESU相同的相关性评估指标，使得这两个级联阶段成为TWIN。因此，CP-GSU始终检索ESU认为重要的item，最大化了行为建模的效果。
• 短期行为建模（Short-term behavior modeling）：从最近的50个行为中提取用户兴趣。该模块关注用户对最近几天的短期兴趣，是TWIN的强有力补充。
• 其他任务建模。除了行为建模，我们还将各种其他任务建模的输出连接起来，包括用户的性别、年龄、职业、位置，视频的持续时间、主题、受欢迎程度、质量，以及播放日期、时间戳、页面位置等上下文特征。

3.3 TWIN: 两阶段兴趣网络

我们将提出的算法命名为TWIN，以突出CP-GSU遵循与ESU相同的相关性评估指标。请注意，这种一致性并不是微不足道的，因为：

• 有效的行为建模算法通常基于多头目标注意力（MHTA）[25]，通过强调目标相关行为来精确捕捉用户兴趣。不幸的是，由于计算复杂度高，MHTA适用的行为序列长度大多限制在几百个之内。
• 为了全面捕捉用户的长期兴趣，CP-GSU应该涵盖最近几个月的用户行为，这可能很容易达到数万个。考虑到在线系统的严格低延迟要求，这个序列长度远远超出了传统MHTA的能力范围。

本节的目的是回答这个关键问题：如何提高MHTA的效率，以便将其从ESU扩展到CP-GSU，或者说从数百个序列长度扩展到至少数万个序列长度？

3.3.1 行为特征分割和线性投影

遵循MHTA [25]的标准符号，我们将长度为𝐿的行为序列$[𝑠_1，𝑠_2，\cdots，𝑠_𝐿]$的特征定义为矩阵𝐾，其中每一行表示一个行为的特征。在实践中，MHTA中注意力得分计算中𝐾的线性投影是阻碍其在极长的用户行为序列上应用的关键计算瓶颈。因此，我们提出以下措施以降低其复杂度。

我们首先将行为特征矩阵𝐾分成两部分：

\(𝐾 ≜ [𝐾_ℎ 𝐾_𝑐] \in R^{𝐿 × (𝐻+𝐶)}\) …(4)

我们将：

• $𝐾_ℎ \in R^{𝐿×𝐻}$：定义为行为items的固有特征（例如视频id、作者、主题、持续时间），它们独立于特定的用户/行为序列
• $𝐾_𝑐 \in R^{𝐿×𝐶}$：定义为user-item交叉特征（例如用户点击时间戳、用户播放时间、点击页面位置、用户-视频交互）。这种分割允许高效计算以下线性投影$𝐾_ℎ 𝑊^ℎ$ $和$𝐾_𝑐 𝑊^𝑐$ 。

对于固有特征$𝐾_ℎ$，虽然维度𝐻很大（每个id特征为64），但线性投影实际上并不昂贵。特定项的固有特征在用户/行为序列之间是共享的。通过必要的缓存策略，𝐾ℎ𝑊 ℎ 可以通过查找和聚集过程高效地“计算”。在线部署的详细信息将在第3.4节介绍。 对于用户-项交叉特征𝐾𝑐，缓存策略不适用，因为：1）交叉特征描述了用户和视频之间的交互细节，因此不在用户行为序列之间共享；2）每个用户最多只观看一次视频。也就是说，在投影交叉特征时没有重复计算。因此，我们通过简化线性投影权重来降低计算成本。

对于用户-项交叉特征$𝐾_𝑐$，缓存策略不适用，因为：

• 1）交叉特征描述了用户和视频之间的交互细节，因此不在用户行为序列之间共享；
• 2）每个用户最多只观看一次视频。也就是说，在投影交叉特征时没有重复计算。

因此，我们通过简化线性投影权重来降低计算成本。

给定𝐽个交叉特征，每个特征的嵌入维度为8（因为没有具有巨大词汇表大小的id特征）。我们将线性投影简化如下:

\(𝐾_𝑐 𝑊^𝑐 ≜ [𝐾_{𝑐,1} w_1^c, \cdots, 𝐾_{𝑐,𝐽} w_𝐽^c]\) … (5)

其中:

• $𝐾_{𝑐,𝑗} \in R^{𝐿×8}$： 是𝐾𝑐的第𝑗个交叉特征的按列切片
• $w_𝑗^c \in R^8$：是其线性投影权重

使用这个简化的投影，我们将每个交叉特征压缩到一个维度，即$𝐾_𝑐 𝑊^𝑐 \in R^{𝐿×𝐽} $ 。请注意，这个简化的投影等价于将$𝑊 _𝑐$ 限制为一个对角块矩阵。

3.3.2 复杂度分析

在传统的MHTA中，线性投影的时间复杂度，即从维度$𝐿×(𝐻+𝐶)$到$𝐿×{d_out}$输出维度的复杂度为𝑂(𝐿×(𝐻+𝐶)×输出维度)。

而在我们的TWIN中的MHTA中，item的固有特征$𝐾_ℎ 𝑊^ℎ$已经预先计算并以𝑂(𝐿)的效率聚合，与维度𝐻无关。而user-item交叉特征$𝐾_𝑐𝑊^𝑐$则被降低为$𝑂(𝐿×𝐶)$的低维计算。 由于𝐶 ≪ 𝐻，且𝐶 ≪ 输出维度，正是这种理论上的加速，使得MHTA在CPGSU和ESU中都能一致地实现。

3.3.3 TWIN中的目标注意力

基于行为的线性投影𝐾ℎ𝑊 ℎ和𝐾𝑐𝑊 𝑐，我们现在定义了目标-行为相关度度量，该度量在CP-GSU和ESU中均匀使用。不失一般性，我们假设用户和目标项之间没有交互，并将目标项的固有特征表示为$q \in R_𝐻$。通过适当的线性投影$𝑊_𝑞$，计算目标项与历史行为之间的相关度分数: $𝜶 ∈ R 𝐿： 𝜶 = (𝐾ℎ𝑊 ℎ ) (q ⊤𝑊 𝑞 ) ⊤ √ 𝑑𝑘

(𝐾𝑐𝑊 𝑐 )𝜷$, (6)

其中$𝑑_𝑘$是查询和键的投影维度。这个相关度分数是通过查询（即目标的固有特征）和键（即行为的固有特征）之间的内积计算的。此外，由于交叉特征被压缩为1维，因此作为偏置项。我们使用𝜷 ∈ R 𝐽作为交叉特征的相对重要性的可学习参数。 在CP-GSU中，这个相关度分数𝜶用于将𝐿 = 104的长期历史行为截断为100个最相关的行为。而在ESU中，我们对最终的100个候选项执行加权平均池化： Attention(q ⊤𝑊 𝑞 , 𝐾ℎ𝑊 ℎ , 𝐾𝑐𝑊 𝑐 , 𝐾𝑊 𝑣 ) = Softmax(𝜶) ⊤𝐾𝑊 𝑣 , (7) 其中𝑊 𝑣是一个投影矩阵。我们稍微滥用了符号，将𝐿 = 100。这个投影𝐾𝑊 𝑣仅在100个行为上执行，因此可以在线高效地进行。我们不需要像计算104个行为的𝜶时那样分割𝐾。 为了共同关注来自不同表示子空间的信息，我们在MHTA中采用了4个头。因此，TWIN的最终输出定义为 TWIN = Concat(head1, …, head4)𝑊 𝑜 , head𝑎 = Attention(q ⊤𝑊 𝑞 𝑎 , 𝐾ℎ𝑊 ℎ 𝑎 , 𝐾𝑐𝑊 𝑐 𝑎 , 𝐾𝑊 𝑣 𝑎 ), 𝑎 ∈ {1, …, 4}, (8) 𝑊 𝑜是一个投影，学习头之间的相对重要性。

略

• https://arxiv.org/pdf/2302.02352.pdf

google youtube搜索团队在《Regression Compatible Listwise Objectives for Calibrated Ranking with Binary Relevance》中提出了一种RCR方法。

摘要

由于LTR（Learning-to-Rank）方法主要旨在提高ranking质量，因此它们的输出分数在设计上并没有进行比例校准（ scale-calibrated）。这从根本上限制了LTR在分数敏感应用（score-sensitive applications）中的使用。虽然有些结合了回归（regression）和排序目标（ranking objective）的简单多目标方法，可以有效地学习比例校准分数（scale-calibrated scores），但我们认为这两个目标不一定兼容，这使得它们之间的权衡不够理想。在本文中，我们提出了一种实用的回归兼容排序（RCR：regression compatible ranking）方法，实现了更好的权衡，其中ranking和regression组件被证明是相互对齐（align）的。虽然同样的思想适用于具有二元（binary）和分级相关性（graded relevance）的排序，但我们在本文中主要关注binary label。我们在几个公共LTR基准测试上评估了所提出的方法，并表明它在回归和排名指标方面始终实现了最佳或有竞争力的结果，并在多目标优化的背景下显著改进了帕累托边界（Pareto frontiers）。此外，我们在YouTube搜索上评估了所提出的方法，并发现它不仅提高了生产环境pCTR模型的ranking质量，还提高了点击预测的准确性。所提出的方法已成功部署在YouTube生产系统中。

1.介绍

LTR（Learning-to-Rank）旨在从训练数据中构建一个排序器（ranker），以便它可以正确地对未见过的对象进行排序。因此，需要ranker在ranking指标（如NDCG）上表现良好。通常情况下，以排序为中心的pairwise或listwise方法（例如RankNet [3]或ListNet [29]）比采用pointwise公式的回归方法实现更好的排序质量。

另一方面，这些应用中的现代系统具有多个阶段，下游阶段会消费前面阶段的预测结果。通常希望ranking分数得到很好的校准，并且分布保持稳定。以在线广告为例，需要对pCTR（预测点击率）模型进行良好的校准，因为它会影响下游拍卖和定价模型[6、16、30]，尽管广告的最终排序对效果来说最为重要。这表明我们希望ranker不仅在排序指标上表现良好，而且在回归指标上也能够将ranker输出分数校准到某个外部尺度上。流行的回归指标：包括用于分级相关性标签（graded relevance labels）的MSE、和用于二元相关性标签（binary relevance labels）的LogLoss。

毫不奇怪，能力强的ranking方法在regression metrics上会表现差些。因为：

• 它们的loss函数对于保序（rank-preserving）的分数变换是不变的，并且倾向于学习未经比例校准的回归目标。
• 这些方法在训练过程中容易出现不稳定，因为所学习的分数可能在连续训练或重新训练中无限发散[30]。

这些因素严重限制了它们在分数敏感应用中的使用。因此，我们别无选择，只能退回到regression-only的方法，即使它们在面向用户的排序指标方面不是最优的。

已经证明，标准的多目标方法可以有效地学习用于ranking的比例校准分数（scale-calibrated scores）[16、25、30、31]。然而，我们认为在这种标准的多目标设置中，regression和ranking目标本质上是相互冲突的，因此最佳权衡可能对其中之一都不理想。在本文中，我们提出了一种实用的回归兼容排序（RCR： regression compatible ranking）方法，其中ranking和regression组件被证明是可以相互对齐的。虽然同样的思想适用于具有二元排序和分级相关性排序，但我们在本文中主要关注二元标签（binary label）。在实证方面，我们在几个公共LTR数据集上进行了实验，并表明所提出的方法在regression和ranking指标方面实现了最佳或竞争结果，并在多目标优化的背景下显著改进了帕累托边界。此外，我们在YouTube搜索上评估了所提出的方法，并发现它不仅提高了生产pCTR模型的ranking能力，还提高了点击预测的准确性。所提出的方法已经在YouTube生产系统中得到了完全部署。

3.背景

学习排序（LTR）关注的问题是：给定一个上下文，学习一个模型来对一个对象列表进行排序。在本文中，我们使用“query”表示上下文，“document”表示对象。在所谓的“打分并排序(score-and-sort)”环境中，学习一个ranker来为每个doc评分，并通过根据分数对docs进行排序来形成最终的ranked list。

更正式地说，假设：

• $𝑞 \in 𝑄$ 为一个query
• $𝑥 \in X$ 为一个doc

则打分函数（score function）定义为：

\[𝑠(𝑞, 𝑥; \theta) : 𝑄 \times X → R\]

其中：

• 𝑄 是query空间
• X 是doc空间
• 𝜽 是打分函数𝑠的参数

一个典型的LTR数据集𝐷由表示为元组$(𝑞, 𝑥, 𝑦) \in 𝐷$的样本组成，其中𝑞，𝑥和𝑦分别为query，doc和label。

假设：

• $q = \lbrace 𝑞 \mid (𝑞, 𝑥, 𝑦) \in 𝐷 \rbrace$：为由𝐷索引的query集合
• $L_{query}(\theta; 𝑞)$：为与单个查询$𝑞 ∈ 𝑄$相关联的loss函数

根据$L_{query}$的定义方式，LTR技术可以大致分为三类： pointwise, pairwise和listwise。

在pointwise方法中，query loss $L_{query}$表示为共享相同query的doc的loss之和。例如，在LR排序（即使用二元相关性标签的排序）中，每个文档的Sigmoid交叉熵损失（用SigmoidCE表示）定义为：

\[SigmoidCE(𝑠, 𝑦) = −𝑦 log \sigma(𝑠) − (1 − 𝑦) log(1 − \sigma(𝑠))\]

…(1)

其中：

• $𝑠 = 𝑠(𝑞, 𝑥; \theta)$：是query-doc pair（𝑞，𝑥）的预测分数
• $\sigma(𝑠) = (1 + exp(−𝑠))−1$：是Sigmoid函数

在文献[30]中表明，SigmoidCE在比例校准方面是可行的，因为当$\sigma(𝑠) \rightarrow E[𝑦 \mid 𝑞, 𝑥]$时，它会达到全局最小值。

在pairwise方法中，query loss $L_{query}$表示为共享相同query的所有doc-doc pair的loss之和。基本的RankNet方法使用pairwise Logistic loss（用PairwiseLogistic表示）[3]：

\[PairwiseLogistic(𝑠_1, 𝑠_2, 𝑦_1, 𝑦_2) = − I(𝑦_2 > 𝑦_1) log \sigma(𝑠_2 − 𝑠_1)\]

…(2)

其中:

• $𝑠_1$和$𝑠_2$是文档$𝑥_1$和$𝑥_2$的预测分数
• I是指示函数
• 𝜎是Sigmoid函数

当$\sigma(𝑠_2 − 𝑠_1) → E[I(𝑦_2 >𝑦_1) \mid 𝑞, 𝑥_1, 𝑥_2]$时，PairwiseLogistic会达到全局最小值，这表明loss函数主要考虑pairwise分数差异，这也被称为平移不变性质（translation-invariant）[30]。

在listwise方法中，query loss $L_{query}$归因于共享相同查询的整个文档列表。流行的ListNet方法使用基于Softmax的Cross Entropy loss（用SoftmaxCE表示）来表示listwise loss[29]：

\[SoftmaxCE(𝑠_{1:𝑁} , 𝑦_{1:𝑁}) = - \frac{1}{C} \sum\limits_{i=1}^N y_i log \frac{exp(s_i)}{\sum\limits_{j=1}^N exp(s_j)}\]

…(3)

其中：

• 𝑁是list size
• $𝑠_𝑖$是预测分数
• $𝐶 = ∑_{𝑗=1}^N 𝑦_𝑗$

在【29】中全局最小值将在以下来实现：

\[\frac{exp(s_i)}{\sum_{j=1}^N exp(s_j)} \rightarrow \frac{E[y_i | q, x_i]}{\sum\limits_{j=1}^N E[y_j | q, x_j]}\]

…(4)

与PairwiseLogistic类似，SoftmaxCE损失是平移不变（translation-invariant）的，并且可能会根据回归指标给出任意更差的分数。

4.REGRESSION COMPATIBLE RANKING

在本节中，我们首先介绍动机，然后正式提出回归兼容排序（RCR）方法。

4.1 动机

文献表明，标准的多目标方法可以有效地学习用于排名的比例校准分数[16、25、30]。以Logistic regression ranking为例，Yan等人将多目标损失定义为SigmoidCE和SoftmaxCE损失的加权和：

\[L_{query}^{MultiObj} (\theta; q) = (1-\alpha) \cdot \sum\limits_{i=1}^N SigmoidCE(s_i, y_i) + \alpha \cdot SoftmaxCE(s_{1:N}, y_{1:N})\]

…(5)

其中：

• 𝛼 ∈ [0, 1]是权衡权重

为简单起见，我们将这种方法称为SigmoidCE + SoftmaxCE。可以看出，SigmoidCE + SoftmaxCE不再是平移不变的（translation-invariant），并且已被证明对于校准排序（calibrated ranking）是有效的。让我们更深入地了解按照这种简单的多目标公式学习的分数是什么。

给定query 𝑞，设$𝑃_𝑖 = E[𝑦_𝑖 \mid 𝑞, 𝑥_𝑖]$为基于在文档$𝑥_𝑖$条件之上的ground truth点击概率。回想一下，当$\sigma(𝑠_𝑖) → 𝑃_𝑖$时，SigmoidCE会达到全局最小值，这意味着对于SigmoidCE，我们会遵循以下的pointwise学习目标：

\[𝑠_𝑖 \rightarrow log 𝑃_𝑖 − log(1 − 𝑃_𝑖)\]

…(6)

另一方面，当以下公式成立时，SoftmaxCE达到全局最小值：

\[\frac{exp(𝑠_𝑖)}{\sum\limits_{𝑗=1}^𝑁 exp(𝑠_𝑗)} \rightarrow \frac{𝑃_𝑖} {\sum\limits_{𝑗=1}^N 𝑃_𝑗}\]

…(7)

或者等价于：

\[𝑠_𝑖 \rightarrow log 𝑃_𝑖 − log \sum\limits_{𝑗=1}^N 𝑃_𝑗 + log \sum\limits_{𝑗=1}^N exp(𝑠_𝑗)\]

…(8)

其中：

• log-∑︁-exp项是未知常数，对最终的SoftmaxCE损失的值或梯度没有影响。

在随机梯度下降的背景下，等式（6）和（8）表明，从SigmoidCE和SoftmaxCE组件生成的梯度将分别将分数推向显著不同的目标。这揭示了标准多目标设置中的两个loss本质上是相互冲突的，将无法找到对两者都理想的解决方案。我们如何解决这个冲突呢？

注意到由于$\sigma(𝑠_𝑖)$在pointwise上趋近于$𝑃_𝑖$，如果我们将等式（8）右侧的ground truth概率$𝑃_𝑖$替换为经验近似项$\sigma(𝑠_𝑖)$，并删除常数项，我们正在构建虚拟的logits：

\[𝑠_𝑖' \leftarrow log \sigma(𝑠_𝑖) − log \sum\limits_{𝑗=1}^N \sigma(𝑠_𝑗)\]

…(9)

如果我们进一步在新的 $logits \ 𝑠_i′$上应用SoftmaxCE loss，我们正在建立以下新的listwise学习目标：

\[\frac{exp(𝑠_𝑖')}{\sum\limits_{𝑗=1}^N exp(s_𝑗')} \rightarrow \frac{𝑃_𝑖}{\sum\limits_{𝑗=1}^N 𝑃_𝑗}\]

…(10)

它等价于：

\[\frac{\sigma(𝑠_𝑖)}{\sum\limits_{𝑗=1}^N \sigma(𝑠_𝑗)} \rightarrow \frac{𝑃_𝑖}{\sum\limits_{𝑗=1}^N 𝑃_𝑗}\]

…(11)

很容易看到，等式（6）自动隐含了等式（11），这意味着，作为pointwise regression和listwise ranking目标，它们在实现全局最小值方面是并行对齐的。

4.2 主方法

受上述示例的启发，我们首先定义一种新的listwise交叉熵损失（ListCE），如下所示。

定义1：设𝑁为列表大小，$𝑠_{1:𝑁}$为预测分数，$𝑦_{1:𝑁}$为label。设$𝑇(𝑠)：R \rightarrow R+$为分数上的非递减变换（non-decreasing transformation）。使用变换𝑇的ListCE定义为：

\[ListCE(𝑇 , 𝑠_{1:𝑁}, 𝑦_{1:𝑁}) = − \frac{1}{𝐶} \sum\limits_{𝑖=1}^N 𝑦_𝑖 log \frac{𝑇(𝑠_𝑖)}{\sum\limits_{𝑗=1}^N 𝑇(𝑠_𝑗)}\]

…(12)

其中：

• $𝐶 = \sum\limits_{𝑗=1}^N 𝑦_𝑗$是一个归一化因子

在本文的范围内，我们可以交替使用带有变换𝑇的ListCE，即ListCE(T)，或者在没有二义性的情况下使用ListCE。我们立即得到以下命题：

命题1：ListCE(exp)简化为SoftmaxCE。

命题2：当满足以下条件时，ListCE(𝑇)可以达到全局最小值：

\[\frac{𝑇(𝑠_𝑖)}{\sum\limits_{𝑗=1}^N 𝑇 (𝑠_𝑗)} \rightarrow \frac{E[𝑦_𝑖 |𝑞, 𝑥_𝑖]}{\sum\limits_{𝑗=1}^N E[𝑦_𝑗 |𝑞, 𝑥_𝑗]}\]

…(13)

证明。设$\bar{𝑦} = E[𝑦

𝑞, 𝑥]$为query-doc对(𝑞, 𝑥)的期望label。在$(𝑥, 𝑦) \in 𝐷$上应用ListCE损失等价于在期望上将其应用于(𝑥,𝑦)。给定变换𝑇和预测分数$𝑠_{1:𝑁}$，其中$𝑝_𝑖 = \frac{𝑇(𝑠_𝑖)} {\sum_{𝑗=1}^N 𝑇(𝑠_𝑗)}$，我们有：

\[ListCE(𝑇 , 𝑠_{1:𝑁}, 𝑦_{1:𝑁}) = \frac{1} {\sum\limits_{𝑗=1}^N 𝑦_𝑗} \sum\limits_{i=1}^N \bar{y_i} log 𝑝_i\]

…(14)

满足：$\sum_{i=1}^N p_i = 1$.

接着构建以下的Lagrangian的公式化：

\[L (𝑝_{1:𝑁}, \lambda) = \frac{1}{\sum\limits_{𝑗=1}^N \bar{𝑦_𝑗}} \sum\limits_{i=1}^N \bar{𝑦_𝑖} log 𝑝_𝑖 + \lambda ( \sum\limits_{i=1}^N 𝑝_𝑖 1)\]

…(15)

找出等式（14）的极值，接着等价于等式（15）的驻点，它满足：

\[\frac{\partial L (𝑝_{1:𝑁}, \lambda)}{\partial 𝑝_𝑖} = \frac{\bar{𝑦_𝑖}}{𝑝_𝑖 \sum\limits_{𝑗=1}^N \bar{𝑦_j}} + \lambda = 0\]

…(16)

并且：

\[\frac{\partial L (𝑝_{1:𝑁}, \lambda)}{\partial \lambda} = \sum\limits_{𝑖=1}^N 𝑝_𝑖 1 = 0\]

…(17)

注意，等式（16）和（17）给出一个在N+1 unknowns上的关于N+1的系统。很容易看到，相等的解决方案是：

\[p_i = \frac{\bar{y_i}}{\sum_{j=1}^N \bar{y_j}}\]

…(18)

并且\(\lambda=1\)。

这意味着唯的的全局极值在：

\[\frac{𝑇 (𝑠_𝑖)}{\sum_{𝑗=1}^N 𝑇(𝑠_𝑗)} \rightarrow \frac{E[𝑦_𝑖 |𝑞, 𝑥_𝑖]}{\sum\limits_{𝑗=1}^N E[𝑦_𝑗 |𝑞, 𝑥_𝑗]}\]

…(19)

很容易验证这个唯一的全局极值归因于全局最小值，这证明了命题。

在逻辑回归排序（logistic-regression ranking）中，所有标签都是二元化的或在[0,1]范围内。一个自然的点对点目标是SigmoidCE损失。使用SigmoidCE作为点对点组件，然后需要使用Sigmoid函数作为变换，以便可以同时进行优化而不产生冲突。 定义2：适用于逻辑回归排名任务（即使用二元相关标签进行排名）中单个查询的回归兼容排名（RCR）损失定义为：

\[L_{query}^{Compatible} (\theta; 𝑞) = (1 − \alpha) \cdot \sum\limits_{𝑖=1}^N SigmoidCE(𝑠_𝑖 , 𝑦_𝑖) + \alpha \cdot ListCE(\sigma, 𝑠_{1:𝑁}, 𝑦_{1:𝑁})\]

…(20)

其中：

• \(\sigma\)是sigmoid funciton

为简单起见，我们将这种方法称为SigmoidCE + ListCE(𝜎)。我们有以下命题：

• 命题3：当$\sigma(𝑠_𝑖) \rightarrow E[𝑦_𝑖

  𝑞, 𝑥_𝑖]$时，SigmoidCE + ListCE(𝜎)可以达到全局最小值。

• 证明：SigmoidCE组件在$\sigma(𝑠_𝑖) \rightarrow E[𝑦_𝑖 \mid 𝑞, 𝑥_𝑖]$时可以达到全局最小值，这意味着：

\[\frac{\sigma(𝑠_𝑖)}{\sum\limits_{𝑗=1}^N \sigma(𝑠_𝑗)} \rightarrow \frac{E[𝑦_𝑖 |𝑞, 𝑥_𝑖]}{\sum\limits_{𝑗=1}^N E[𝑦_𝑗 |𝑞, 𝑥_𝑗]}\]

…(21)

它会最小化ListCE(𝜎)在它的全局最小值上。

#

略

• https://arxiv.org/pdf/2211.01494.pdf

介绍

字节在《Monolith: Real Time Recommendation System With Collisionless Embedding Table》提出了它们的embedding table实现。

摘要

对于许多依赖于时间敏感客户反馈的业务来说，构建一个可扩展且实时的推荐系统至关重要，例如短视频排序或在线广告。尽管像TensorFlow或PyTorch这样的生产规模深度学习框架被广泛采用，但这些通用框架在推荐场景中的业务需求方面存在多种不足：

• 一方面，基于静态参数和dense计算调整系统对于具有动态和稀疏特征的推荐是不利的；
• 另一方面，这些框架设计时将批量训练阶段和服务阶段完全分离，阻止了模型与客户反馈实时互动。

这些问题促使我们重新审视传统方法并探索根本不同的设计选择。在本文中，我们介绍了Monolith1，一个为在线训练量身定制的系统。我们的设计理念受到了我们的应用工作负载和生产环境的观察，这与其他推荐系统有明显的不同。我们的贡献是多方面的：

• 首先，我们制作了一个无冲突的嵌入表，并进行了诸如可过期嵌入和频率过滤等优化以减少其内存占用；
• 其次，我们提供了一个具有高容错性的生产就绪在线训练架构；
• 最后，我们证明了系统可靠性可以与实时学习进行权衡。

Monolith已成功应用于BytePlus Recommend2产品中。

1 引言

过去十年见证了由推荐技术驱动的业务的蓬勃发展。为了追求更好的用户体验，为每个用户实时提供个性化内容是这些商业应用的共同目标。为此，用户最新互动的信息通常被用作训练模型的主要输入，因为它能最好地描绘用户画像，并预测用户的兴趣和未来行为。

深度学习已经在推荐模型中占据主导地位[5, 6, 10, 12, 20, 21]，因为海量的用户数据天然适合大规模数据驱动的神经网络模型。然而，在工业级推荐系统中利用深度学习的力量，不断遇到由现实世界用户行为数据的独特特性引发的问题。这些数据在两个方面与用于传统深度学习问题（如语言建模或计算机视觉）的数据截然不同：

• (1) 特征大多是稀疏的、类别型的并且动态变化的；
• (2) 训练数据的底层分布是非平稳的，即概念漂移（Concept Drift）[8]。

这些差异给从事推荐系统的研究人员和工程师带来了独特的挑战。

1.1 稀疏性和动态性

推荐的数据大多包含稀疏的类别型特征（sparse categorical features），其中一些特征出现的频率很低。将它们映射到高维嵌入空间的常见做法会引发一系列问题：

• 与单词片段数量有限的语言模型不同，推荐系统中的用户和ranking items的量级要大得多。如此庞大的嵌入表几乎无法适应单个主机内存；
• 更糟糕的是，随着更多用户和item的加入，嵌入表的大小预计会随着时间增长，而像[1, 17]这样的框架使用固定大小的dense变量来表示嵌入表。

在实践中，许多系统采用低冲突哈希[3, 6]来减少内存占用，并允许ID的增长。这依赖于一个过于理想化的假设，即嵌入表中的ID频率分布均匀，并且冲突对模型质量无害。不幸的是，这对于现实世界的推荐系统很少是真的，其中一小部分用户或item的出现次数明显更多。随着嵌入表大小的自然增长，哈希键冲突的几率增加，导致模型质量恶化[3]。

因此，对于生产规模的推荐系统来说，自然需要有能力在其参数中捕获尽可能多的特征，并且还要有能力灵活调整它试图管理的用户和item的数量。

1.2 非平稳分布

视觉和语言模式在几个世纪的时间尺度上几乎不会发展，而对一个话题感兴趣的用户可能在下一分钟就转移他们的热情。因此，用户数据的底层分布是非平稳的，这种现象通常被称为概念漂移[8]。

直观地说，更近期的历史信息可以更有效地预测用户行为的变化。为了减轻概念漂移的影响，服务模型需要尽可能接近实时地从新的用户反馈中更新，以反映用户的最新兴趣。

鉴于这些区别，并观察到我们生产环境中出现的问题，我们设计了一个大规模推荐系统Monolith来解决这些痛点。我们进行了广泛的实验来验证和迭代我们的设计。Monolith能够：

• (1) 通过设计一个无冲突的哈希表和一个动态特征淘汰机制，为稀疏特征提供完整的表达能力；
• (2) 通过在线训练，将服务反馈实时循环回训练。

凭借这些架构能力，Monolith在大约相似的内存使用情况下，始终优于采用有冲突的哈希技巧的系统，并实现了最先进的在线服务AUC，而没有过度负担我们服务器的计算能力。

本文的其余部分组织如下。我们首先在第2节详细阐述Monolith如何通过无冲突哈希表和实时训练解决现有挑战的设计细节。第3节将展示实验结果，以及生产测试的结论和对时效性、可靠性和模型质量之间权衡的一些讨论。第4节总结相关工作并与Monolith进行比较。第5节结束本文。

2 设计

Monolith的整体架构通常遵循TensorFlow的分布式Worker-ParameterServer设置（图2）。在Worker-PS架构中，机器被分配不同的角色；Worker机器负责执行图定义的计算，而PS机器存储参数并根据Worker计算的梯度更新它们。

图2 Worker-PS架构

在推荐模型中，参数被分为两组：dense和sparse：

• dense参数是深度神经网络中的权重/变量
• sparse参数指的是对应稀疏特征的嵌入表

在我们的设计中，dense和sparse参数都是TensorFlow图的一部分，并存储在参数服务器上。

与TensorFlow的密集参数变量类似，我们为稀疏参数设计了一套高效、无冲突且灵活的哈希表操作。作为补充TensorFlow训练和推理分离限制的Monolith，其弹性可扩展的在线训练旨在在短间隔内高效地将参数从【训练-PS】同步到【在线服务-PS】，模型的鲁棒性由容错机制保证。

2.1 哈希表

我们在设计sparse参数表示时的一个首要原则是：避免将不同ID的信息压缩到同一固定大小的嵌入中。使用现成的TensorFlow变量模拟动态大小的嵌入表不可避免地会导致ID冲突，随着新ID的到来和表的增长，这种情况会加剧。

因此，我们没有在变量的基础上构建，而是为我们的sparse参数开发了一个新的键值哈希表。

我们的哈希表在底层使用Cuckoo哈希图[16]，它支持插入新键而不与现有键冲突。Cuckoo哈希在查找和删除上实现了最坏情况下的𝑂(1)时间复杂度，以及预期的平均𝑂(1)时间复杂度的插入。如图3所示，它维护两个表$𝑇_0,𝑇_1$，具有不同的哈希函数$ℎ_0(𝑥), ℎ_1(𝑥)$，一个元素将被存储在它们中的一个。当尝试将元素𝐴插入$𝑇_0$时，它首先尝试将𝐴放置在$ℎ_0(𝐴)$；如果$ℎ_0(𝐴)$被另一个元素𝐵占用，它会将𝐵从$𝑇_0$中驱逐出去，并尝试使用相同的逻辑将𝐵插入$𝑇_1$。这个过程将重复进行，直到所有元素稳定，或者在插入遇到循环时发生重新哈希。

图3 布谷鸟哈希（Cuckoo HashMap）

在我们的设计中，内存占用减少也是一个重要考虑因素。简单地将每个新ID插入哈希表会迅速耗尽内存。对真实生产模型的观察导致两个结论：

• (1) 只出现几次的ID对提高模型质量的贡献有限。一个重要的观察是，ID是长尾分布的，其中流行的ID可能出现数百万次，而不受欢迎的ID出现不超过十次。对应这些不频繁ID的嵌入由于缺乏训练数据而拟合不足，模型将无法基于它们做出良好的估计。归根结底，这些ID不太可能影响结果，因此从这些低频ID中移除不会影响模型质量；
• (2) 来自遥远历史的陈旧ID很少对当前模型做出贡献，因为它们中的许多从未被访问过。这可能是因为一个不再活跃的用户，或者一个过时的短视频。存储这些ID的嵌入对模型没有任何帮助，只会白白消耗我们的PS内存。

基于这些观察，我们为哈希表设计了几项特征ID过滤启发式方法，以实现更内存高效的实现：

• (1) 在ID被允许进入嵌入表之前进行过滤。我们有两种过滤方法：首先，我们根据它们出现的次数在它们被插入为键之前进行过滤，出现次数的阈值是一个可调的超参数，每个模型各不相同；此外，我们使用概率过滤器进一步减少内存使用；
• (2) ID被定时，并在一定时间内不活跃后被设置为过期。过期时间也是每个嵌入表可调的，以允许区分对历史信息敏感度不同的特征。

在我们的实现中，哈希表被实现为TensorFlow资源操作。与变量类似，查找和更新也被实现为原生TensorFlow操作，以便于集成和更好的兼容性。

2.2 在线训练

在Monolith中，训练被分为两个阶段（图1）：

图1 Monolith在线训练架构

• (1) 批量训练阶段。这个阶段作为一个普通的TensorFlow训练循环工作：在每个训练步骤中，训练工作器从存储中读取一小批训练样本，从PS请求参数，计算前向和反向传播，最后将更新后的参数推送到training PS。与其他常见的深度学习任务略有不同，我们只对数据集进行一次遍历的训练。批量训练对于我们在修改模型架构并重新训练模型时训练历史数据很有用；
• (2) 在线训练阶段。模型部署到在线服务后，训练不会停止，而是进入在线训练阶段。训练工作器不是从存储中读取小批量样本，而是实时消费实时数据并更新training PS。training PS定期将参数同步到serving PS，这将立即在用户端生效。这使我们的模型能够根据用户的反馈实时互动适应。

2.2.1 流式引擎

Monolith构建了无缝切换批量训练和在线训练的能力。这是通过我们设计的流式引擎实现的，如图4所示。 在我们的设计中，我们使用一个Kafka队列来记录用户的行为（例如点击一个项目或喜欢一个项目等），另一个Kafka队列用于特征。引擎的核心是一个Flink流式作业在线特征Joiner。在线Joiner将特征与用户行为的标签连接起来，生成训练样本，然后写入Kafka队列。训练样本队列被在线训练和批量训练都消费：

图4 Streaming Engine

• 对于在线训练，训练工作器直接从Kafka队列读取数据；
• 对于批量训练，数据转储作业首先将数据转储到HDFS；在HDFS中累积了一定量的数据后，训练工作器将从HDFS检索数据并执行批量训练。

training PS中更新的参数将根据参数同步计划推送到serving PS。

2.2.2 在线Joiner

在现实世界的应用中，用户行为日志和特征是无时间顺序保证地流式传输到在线Joiner（图5）。因此我们使用每个请求的唯一键，以便用户行为和特征能够正确配对。用户行为的延迟也可能是一个问题。例如，用户可能在几天前他们被展示的项目后决定购买。这对于Joiner来说是一个挑战，因为如果所有特征都保留在缓存中，它将无法适应内存。在我们的系统中，使用磁盘上的键值存储来存储等待超过一定时间周期的特征。当用户行为日志到达时，它首先查找内存缓存，如果缓存缺失，则查找键值存储。

图5 Online Joiner

在现实世界的应用中出现的另一个问题是，负样本和正样本的分布高度不均匀，前者的数量可能比后者高几个数量级。为了防止正样本被负样本淹没，一个常见的策略是进行负采样。这肯定会改变训练模型的底层分布，将其调整为更高概率的正预测。作为补救措施，我们在服务期间应用对数几率校正[19]，确保在线模型是原始分布的无偏估计器。

2.2.3 参数同步

在在线训练期间，Monolith训练集群不断从在线服务模块接收数据并更新training PS上的参数。使在线serving PS能够从这些新训练的参数中受益的一个关键步骤是：同步更新的模型参数。在生产环境中，我们遇到了几个挑战：

• 在线serving PS上的模型在更新时不能停止服务。我们生产中的模型通常有数TB的大小，因此替换所有参数需要一段时间。在替换过程中停止在线PS服务模型是不可接受的，更新必须即时进行；
• 从training PS到在线serving PS传输数TB的模型将对网络带宽和PS上的内存造成巨大压力，因为这需要双倍的模型大小内存来接受新到达的模型。

为了使在线训练能够扩展到我们业务场景的规模，我们设计了一种增量式的即时定期参数同步机制，基于我们模型的几个显著特征：

• (1) sparse参数主导了推荐模型的大小；
• (2) 给定一个短时间窗口，只有一小部分ID会被训练，它们的embedding会被更新；
• (3) dense变量的变动速度远慢于sparse嵌入。这是因为：在基于动量的优化器（momentum-based optimizers）中，dense变量的动量积累被推荐训练数据的庞大size所放大，而单个数据批次中只有少数sparse嵌入接收更新。

(1) 和 (2) 允许我们利用所有特征ID的sparse更新。在Monolith中，我们维护一个被触摸键（touched keys）的哈希集合，代表自上次参数同步以来embedding中被训练的ID。我们以分钟级别的时间间隔将被触摸键集中的稀疏参数子集从training PS推送到在线serving PS。这种相对较小的增量参数更新包对网络传输来说很轻，并且在同步过程中不会导致内存急剧增加。

我们还利用 (3) 进一步减少网络I/O和内存使用，通过为稀疏参数设置更积极的同步计划，而不太频繁地更新密集参数。这可能会导致我们服务的dense参数与sparse部分相比是相对陈旧的版本。然而，由于 (3) 中提到的原因，这种不一致是可以容忍的，因为没有观察到明显的损失。

图6 DeepFM架构

2.3 容错性

作为一个生产系统中的系统，Monolith被设计为在PS（Parameter Server）失败时能够恢复。容错的一个常见选择是：定期对模型的状态进行快照，并在检测到PS故障时从最新的快照中恢复。快照频率的选择有两个主要影响：

• (1) 模型质量。直观上，随着快照频率的增加，模型质量受到近期历史丢失的影响较小。
• (2) 计算开销。对多TB模型进行快照并非没有成本。它会产生大量的内存复制和磁盘I/O。

作为模型质量和计算开销之间的权衡，Monolith每天都会对所有training PS进行快照。尽管在故障情况下PS会丢失一天的更新，但我们通过实验发现性能下降是可以接受的。我们将在下一节分析PS可靠性的影响。

3 评估

为了更好地理解我们提出的设计带来的益处和权衡，我们在生产规模上进行了一系列实验，并使用实时服务流量进行了A/B测试，以从不同方面评估和验证Monolith。我们希望通过实验回答以下问题：

• (1) 无冲突哈希表能带来多少好处？
• (2) 实时在线训练有多重要？
• (3) 在大规模生产场景中，Monolith的参数同步设计是否足够健壮？

在本节中，我们首先介绍我们的实验设置，然后详细讨论结果和我们的发现。

3.1 实验设置

3.1.1 嵌入表

如第2.1节所述，Monolith中的嵌入表实现为无冲突哈希表。为了证明避免嵌入表中冲突的必要性并量化我们无冲突实现的收益，我们在Movielens数据集和我们的内部生产数据集上分别进行了两组实验：

(1) MovieLens ml-25m数据集[11]。这是一个标准的公共电影评分数据集，包含约2500万个评分，涉及约162000名用户和62000部电影。

• 标签预处理。原始标签是0.5到5.0的评分，而在生产中我们的任务大多是接收用户的二元信号。为了更好地模拟我们的生产模型，我们将刻度标签转换为二元标签。

(2) 内部推荐数据集。我们还在生产环境中的推荐模型上进行了实验。这个模型通常遵循多塔架构，每个塔负责学习预测一种专门的用户行为。

• 每个模型大约有1000个嵌入表，嵌入表的大小分布非常不均匀；
• 嵌入表的原始ID空间是$2^{48}$。在我们的基线中，我们应用了一种哈希技巧，通过分解来限制嵌入表的大小。具体来说，我们使用两个较小的嵌入表而不是一个巨大的表来为每个ID生成一个唯一的嵌入，通过向量组合：

\[ID_r = ID\ \% \ 2^{24} \\ ID_q = ID\ / \ 2^{24} \\ E = \mathbf{E}_{\text{l}} + \mathbf{E}_{\text{q}}\]

其中：

• $E_l$, $E_q$ ：分别对应于 $I_l$, $I_q$ 的嵌入。

这有效地将嵌入表的大小从$2^{48}$减少到$2^{25}$；

• 这个模型正在实时生产中服务，这个实验的性能是通过在线AUC和实时服务流量来衡量的。

3.1.2 在线训练

在在线训练期间，我们以分钟级别的间隔用最新的参数集更新我们的在线serving PS。我们设计了两组实验来验证模型质量和系统鲁棒性。

(1) 更新频率。为了调查分钟级更新频率的必要性，我们进行了实验，以不同的间隔从训练模型同步参数到预测模型。

我们使用的是Criteo Display Ads Challenge数据集，这是一个大规模的标准数据集，用于基准测试CTR模型。它包含了7天按时间顺序排列的数据记录特征和点击行为。在这个实验中，我们使用了一个标准的DeepFM模型，如第6节所述。为了模拟在线训练，我们对数据集进行了以下预处理。我们从数据集中取出7天的数据，并将其分为两部分：5天的数据用于批量训练，2天的数据用于在线训练。我们进一步将2天的数据按时间顺序分成N个片段。在线训练通过算法1模拟。因此，我们模拟了以数据片段数量确定的时间间隔将训练参数同步到在线serving PS的过程。我们尝试了N = 10, 50, 100，大致对应于5小时、1小时和30分钟的更新间隔。

算法1

（2）实时实验。此外，我们还进行了一个实时实验，使用真实的服务流量进一步展示在线训练在现实世界应用中的重要性。这个A/B实验比较了我们的一个生产广告模型中的在线训练和批量训练。

3.2 结果和分析

3.2.1 嵌入冲突的影响

来自MovieLens数据集和内部推荐数据集的结果都显示，嵌入冲突会危及模型质量。

图7 DeepFM模型在MovieLens数据集上的embedding冲突的效果

（1）无冲突哈希表的模型始终优于有冲突的模型。这一结论无论在以下情况下都成立：

• 训练周期数量的增加。如图7所示，无冲突嵌入表的模型从第一个周期开始就有更高的AUC，并在更高的值处收敛；
• 由于概念漂移，分布随时间的变化。如图8所示，无冲突嵌入表的模型也随着时间的推移和用户/项目上下文的变化而保持稳健。

图8 在生产环境下推荐模型的embedding冲突效果

（2）由无冲突嵌入表引起的数据稀疏性不会导致模型过拟合。如图7所示，无冲突嵌入表的模型在收敛后不会过拟合。

3.2.2 在线训练：

实时性与可靠性的权衡。我们发现，更高的参数同步频率总是有助于提高在线服务AUC，并且在线服务模型对PS（Parameter Server）部分数据丢失的容忍度超出我们的预期。

（1）参数同步频率的影响。在我们使用Criteo Display Ads Challenge数据集进行的在线流式训练实验中，模型质量随着参数同步频率的增加而持续提高，这可以从两个角度明显看出：

图9 在Criteo数据集上Online training vs. Batch training，蓝线：online training模型的AUC；黄线：batch training模型的AUC

• 进行在线训练的模型比没有进行在线训练的模型表现更好。图9a、9b、9c比较了在线训练模型按后续数据片段评估的AUC与批量训练模型按每个数据片段评估的AUC；
• 参数同步间隔较小的模型比间隔较大的模型表现更好。图10和表2比较了同步间隔为5小时、1小时和30分钟的模型的在线服务AUC。

图10 online training中不同同步间隔的比较

在生产环境中，在线训练与批量训练的实时A/B实验也显示在线服务AUC有显著提升（表3）。

受此观察启发，我们将稀疏参数尽可能频繁地同步到生产模型的serving PS（目前是分钟级），以忍受计算开销和系统可靠性的程度。回想第2.2.3节中提到的密集变量需要较不频繁的更新，我们每天更新它们。这样做，我们可以将计算开销降到非常低的水平。假设每分钟有100,000个ID更新，嵌入的维度是1024，需要传输的总数据大小是4KB × 100,000 ≈ 400MB每分钟。

对于密集参数，由于它们是每天同步的，我们选择在流量最低的时候（例如午夜）安排同步。（2）PS可靠性的影响。在分钟级参数同步的情况下，我们最初期望更频繁地对training PS进行快照以匹配实时更新。令人惊讶的是，我们将快照间隔扩大到1天，仍然几乎观察不到模型质量的损失。

在个性化排序系统中，找到模型质量和计算开销之间的正确权衡是困难的，因为用户对推荐质量非常敏感。传统上，大规模系统倾向于为它们的模型设置频繁的快照计划，以牺牲计算资源为代价，以最小化模型质量的损失。我们也在这方面做了很多探索，令人惊讶的是，模型质量比预期的更稳健。在PS机器每天有0.01%的故障率的情况下，我们发现前一天的模型出奇地好用。这个可以通过以下计算来解释：假设一个模型的参数分布在1000个PS上，并且它们每天快照一次。鉴于0.01%的故障率，每10天就会有其中一个故障，我们失去了这个PS上一天的所有更新。假设日活跃用户（DAU）为1500万，用户ID在每个PS上均匀分布，我们每10天就会失去来自15000用户的一天反馈。这是可以接受的，因为：

-（a）对于用户特定的稀疏特征，这相当于失去了0.01% DAU的微小部分； -（b）对于密集变量，由于它们更新缓慢，如我们在2.2.3节中讨论的，失去1000个PS中一天的更新是微不足道的。

基于上述观察和计算，我们大幅降低了快照频率，从而节省了大量的计算开销。

4 相关工作

自从深度学习在工业级推荐系统中最早成功应用以来[6, 10]，研究人员和工程师一直在采用各种技术来改善第1节中提到的问题。

为了解决稀疏特征表示的问题，[3, 6]使用固定大小的嵌入表和哈希技巧。还有尝试改进哈希以减少冲突[3, 7]。其他工作直接使用原生键值哈希表，以允许表大小的动态增长[12, 15, 20, 21]。这些实现基于TensorFlow，但依赖于特别设计软件机制[14, 15, 20]或硬件[21]来访问和管理它们的哈希表。与这些解决方案相比，Monolith的哈希表是另一种原生TensorFlow操作。它对开发者友好，具有更高的跨平台互操作性，适合ToB场景。与TensorFlow的有机紧密集成还使得计算性能的优化更容易。

弥补训练和部署之间的差距和缓解概念漂移[8]是另一个感兴趣的话题。为了支持在线更新并避免内存问题，[12]和[20]设计了特征逐出机制，以灵活调整嵌入表的大小。[12]和[14]都支持某种形式的在线训练，其中学习到的参数与传统批量训练相比，以相对较短的时间间隔同步到服务，具有容错机制。Monolith采取了类似的方法来弹性地接纳和逐出特征，同时它有一个更轻量级的参数同步机制来保证模型质量。

参考

• 1.https://arxiv.org/pdf/2209.07663

介绍

Nvidia在《A GPU-specialized Inference Parameter Server for Large-Scale Deep Recommendation Models》提出了面向大规模深度推荐模型的GPU专用推理参数服务器。

摘要

推荐系统对于各种现代应用和网络服务至关重要，例如新闻推送、社交网络、电子商务、搜索等。为了实现最高的预测准确性，现代推荐模型结合深度学习和数万亿级别（terabyte）的嵌入表，以获得底层数据的细粒度表示。传统的推理服务架构需要将整个模型部署到独立的服务器上，但在这种大规模下这是不可行的。

在本文中，我们提供了关于在线推荐系统的有趣且具有挑战性的推理领域的见解。我们提出了HugeCTR分层参数服务器（HPS:Hierarchical Parameter Server），一种行业领先的分布式推荐推理框架，结合了高性能GPU嵌入缓存和分层存储架构，以实现在线模型推理任务的低延迟检索嵌入。HPS的特点包括：

• （1）冗余的分层存储系统；
• （2）一种新颖的高带宽缓存，以加速NVIDIA GPU上的并行嵌入查找；
• （3）在线训练支持；
• （4）轻量级API，便于轻松集成到现有的大规模推荐工作流程中。

为了展示其能力，我们使用合成工程和公共数据集进行了广泛的研究。我们发现，HPS可以显著降低端到端推理延迟，根据batch-size大小，相较于CPU基线实现，为流行的推荐模型实现了5~62倍的加速。通过多GPU并发部署，HPS还可以大大提高推理QPS。

1 引言

推荐系统（RS）广泛应用于各种应用和在线服务中，例如新闻推送、电子商务、社交网络、搜索等。为了提供准确的预测，最先进的算法依赖于基于embedding的深度学习模型。图1展示了一个深度推荐模型（DLRM）的典型架构。输入包括dense特征（例如，年龄、价格等）和sparse特征（例如，用户ID、类别ID等）。sparse特征通过在嵌入表中查找转换为dense的embedding向量，以便将这些向量与dense特征结合后输入到一些dense连接的深度学习模型（例如，多层感知机MLP、Transformer等[38, 39]）中，从而预测点击率（CTR）。

图1

embedding可能会占用数据中心内存容量的很大一部分。通常，从集中式参数服务器中检索这些embedding会花费大量时间，这增加了延迟，从而拖慢了后续的计算。与面向吞吐量的训练系统[5, 7, 12–14, 16, 17, 22, 42]不同，在线推理系统严格受限于延迟要求[40]。因此，embedding查找速度对于深度推荐模型的推理性能至关重要。

在推理过程中，每个小batch的数据通常引用数万个embedding。通过键值对每个embedding进行穷举搜索，需要参数服务器遍历某些内部数据结构。从嵌入表中查找单个embedding通常是独立的，因此很容易并行化。同时，现代GPU架构允许调度数千个线程并发运行，其内存子系统采用了特殊的内存技术，提供比同等CPU内存更高的带宽和吞吐量[27]。这些特性使得GPU架构非常适合处理embedding向量查找的工作负载。

挑战。最先进的推荐模型中使用的嵌入表规模可能非常庞大，通常从数十GB到数TB不等，这远远超出了大多数GPU的内存容量。此外，在线推理期间的batch-size通常太小，无法有效利用单个GPU的大规模并行处理优化的计算资源。因此，嵌入查找工作负载需要大量的GPU内存，但只需要很少的计算资源。这种需求的不平衡与现有硬件显著偏离，降低了GPU在推理系统中的吸引力。因此，大多数现有解决方案将嵌入查找操作与dense计算（即模型的其余部分）解耦，后者在GPU中执行，而将嵌入查找操作移至CPU[21]。这样一来，它们放弃了GPU的内存带宽优势，而CPU以及CPU与GPU之间的通信带宽成为主要瓶颈。结果，GPU的不成比例的处理能力在这种设置中大多处于闲置状态（=资源浪费）。

方法。通常无法将所有嵌入表完全保留在GPU内存中。然而，实际推荐数据集的实证研究表明，在CTR和其他推荐任务的推理过程中，嵌入键(embedding key)的访问通常表现出很强的局部性，并且大致遵循幂律分布[5, 7, 12, 17]。因此，每个小batch中的大部分嵌入键仅引用一小部分热门嵌入。将这些热门嵌入缓存在GPU内存中（模型的其余部分也在GPU中处理），可以实现部分GPU加速的嵌入查找。基于这些观察，我们构建了一个推理框架，即HugeCTR分层参数服务器（HPS），以充分利用GPU资源，同时不受GPU内存限制的约束。特别是，HPS引入了一种GPU嵌入缓存数据结构，试图将热门嵌入保留在GPU内存中。该缓存由一个参数服务器补充，该服务器保存所有嵌入表的完整副本。我们的贡献可以总结如下：

• 分层数据库架构：允许利用集群内存资源，并提供异步更新机制，以在在线推理期间保持较高的GPU嵌入缓存命中率。
• 高性能动态GPU嵌入缓存：通过跟踪并缓存高频出现的嵌入到高吞吐量的GPU内存中，同时重叠主机/设备数据传输，从而最大化吞吐量。
• 在线模型更新机制：支持分布式推理部署（即实时更新）。
• 可定制的HPS后端：为NVIDIA Triton GPU推理服务器[31]提供并发模型执行、混合模型部署和集成模型管道服务。

本文的结构如下。在第2节中，我们对支撑我们方法的核心概念进行了基础性讨论。随后，在第3至第5节中，我们介绍并讨论了HPS的各个组件及其交互方式。在第6节中，我们讨论了HPS如何实现实时模型更新。最后，在第7节中，我们通过实验研究评估了HPS的性能，并在第8节中提供了结论性意见。

2 背景

2.1 嵌入表

当前广告、推荐和搜索领域的主流算法采用了一种将嵌入表与深度神经网络结合的模型结构，形成深度学习推荐模型（DLRM）[24]。这类模型的基础是嵌入 $ e $，它表示用户或物品特征的学习数值表示，以dense向量的形式在某个 $ d $ 维空间中对齐（$ e \in \mathbb{R}^d $）。我们用 $ E_j = {e^0_j, e^1_j, \dots, e^n_j} $ 表示某个特征 $ j $ 的嵌入的离散子集。为了在模型中方便访问，我们将这些嵌入组织为嵌入特征表，形式如下：

\[T_j = \langle K_j, E_j \rangle = \{\langle k^0_j, e^0_j \rangle, \langle k^1_j, e^1_j \rangle, \dots, \langle k^n_j, e^n_j \rangle\}, \quad (1)\]

其中：

• 每个元组 $ \langle k^i_j, e^i_j \rangle $ 包含一个键 $ k^i_j $，用于标识和引用第 $ i $ 个嵌入表条目 $ e^i_j $。
• 键空间 $ K_j = {k^0_j, k^1_j, \dots, k^n_j} $ 是离散的，且满足 $ \forall k^i_j, k^z_j \in K_j $ （$ i \neq z \rightarrow k^i_j \neq k^z_j $）。每个键的值取决于底层数据或任务。通常，键空间是稀疏的。

为了评估DLRM的CTR（参见图1），驱动应用程序必须首先从嵌入表中选择与预测相关的条目。这可以通过从每个嵌入特征表的查询键子集 $ Q_j $ 中查找键来完成（即 $ Q = {Q_0 \subseteq K_0, Q_1 \subseteq K_1, \dots} $）。因此，$ Q_j = {q^0j, q^1_j, \dots, q^m_j} $ 表示从 $ T_j $ 中查找 $ m $ 个对应嵌入条目的查询。对应的结果集为 $ R{Q_j} = {q^0_j \mapsto e^0_j, q^1_j \mapsto e^1_j, \dots, q^m_j \mapsto e^m_j} $。我们的主要目标是加速大规模检索此类结果集。

2.2 去重与偏斜性

为了避免在多次需要相同嵌入表条目时进行不必要的重复查找，HugeCTR 在执行任何后续步骤之前始终应用去重操作（即 $ Q^* = \text{dedup}(Q) $）。这对于小batch处理尤为重要，其中： $ Q $ 是许多输入样本的拼接。自然，如果查询分布 $ Q $ 的偏斜性增加，去重的效果会更加显著。

理解并利用数据集的偏斜特性对于实现峰值效率至关重要。许多现实世界中的推荐数据集（例如 Criteo [6]）呈现出幂律分布 [3]。也就是说，某些键子集比其他键更频繁地被引用，使得从 $ Q_j $ 中采样 $ q_j $ 最终近似满足 $ p(x) \propto x^{-\alpha} $。图 2 展示了一个场景，其中嵌入键的召回统计量近似于幂律分布。

图2

键空间可以分为三类：

1. 频繁嵌入：这些嵌入几乎出现在每个batch中，占据了召回/更新请求的很大一部分。频繁集通常很小，即使对于大型嵌入语料库，也只有几千个嵌入会如此频繁地出现。
2. 随机嵌入：这些嵌入每隔几个batch出现一次（即随着时间的推移相对规律地出现）。
3. 稀有嵌入：这些嵌入位于分布的另一端，在查询中出现的频率非常低。

由于请求会反复引用频繁和随机嵌入，对它们应用高效的缓存方法可以最大程度地提高整体系统性能。我们的 HPS 设计（见第 3 节）正是基于这一观察。

如果查询数据集是固定的，嵌入的类别划分是绝对确定的。在训练 HugeCTR 模型时，我们利用这一点来实现世界级的模型收敛速度 [8, 19]。在在线推理过程中，召回统计量取决于实际传入的用户请求，这些请求无法被预先预测。由于突发事件、趋势或时尚的变化，单个嵌入的类别划分可能会随时间而变化。对于大多数推荐任务而言，运行时统计量处于不断变化中。因此，推理系统必须具备自适应性。

2.3 GPU加速推理架构

用于机器学习推理工作负载的参数服务器，主要依赖于可以轻松与GPU并行化的数据库操作[35, 37, 41]。需要快速响应时间的应用，例如在线事务处理（OLTP），通常能够从GPU加速中受益匪浅[1]。然而，GPU内存限制带来了严峻的挑战。为了实现可扩展性，许多现有的GPU加速数据库系统以及我们的方法都采用了分层存储架构，通过其他存储资源扩展可用的GPU内存。由于外部内存资源的访问效率无法与本地GPU内存相媲美[27]，因此在这些系统中，与主机系统的数据交换性能被特别强调[23]。为了实现峰值性能，必须将重叠查询处理与高效的通信模式和数据放置策略结合使用，并在运行时动态优化[1, 2, 20]。

为机器学习平台构建参数服务器面临许多挑战[2, 5, 7, 12–14, 16, 17, 20, 22, 42]。在设计用于推理生产环境的混合GPU/CPU架构时，至少需要克服两个主要瓶颈：

1. 高延迟：由于CPU和GPU之间通信时的DRAM带宽限制[18, 40]。
2. 部署延迟：由于在线训练导致的模型规模和复杂性增加，因为快速的增量模型更新在数据一致性和带宽方面提出了巨大挑战。

为了解决这些瓶颈，我们的HPS专门设计为用于大规模推荐模型的GPU推理参数服务器。它处理数据同步和通信，以在不同推理节点之间共享模型参数（嵌入表）[26]，并执行各种优化以提高并行多模型/多GPU推理期间的GPU利用率，包括将分布式嵌入表组织为分区[25]、GPU友好的缓存[30]以及异步数据移动机制[29]。

3 分层参数服务器

我们的分层参数服务器（HPS）使HugeCTR能够使用具有巨大嵌入表的模型进行推理。这是通过利用集群中的CPU内存资源扩展嵌入存储空间，超越GPU的限制来实现的。HPS的设计目标是解决传统CPU参数服务器方法通常面临的三个主要挑战：

1. 模型参数的下载/流式传输：从CPU内存中集中维护的嵌入表分区到各个GPU计算设备上的模型实例。如果嵌入表无法完全装入GPU内存，这个问题会被放大。HPS通过利用数据分布的局部性的GPU缓存机制，极大地缓解了这一问题。

2. 因推理平台高可用性需求和带宽限制而增加的部署成本：通过联合组织和使用推理集群的分布式CPU内存，HPS节省了资源，并实现了即时在线模型更新（即从训练到推理的更新）。

3. GPU缓存与参数服务器之间的参数更新和刷新：如果仅将模型的一部分加载到GPU内存中，则在查找期间GPU上可能会遗漏参数，这尤其具有挑战性。HPS通过异步插入和刷新机制处理CPU和GPU之间的额外参数交换，以保持参数的一致性。

3.1 存储架构

我们的HPS实现为一个3级分层内存架构（参见图3），利用GPU的GDDR和/或高带宽内存（HBM）、分布式CPU内存以及本地SSD存储资源。

图3

这些组件之间的通信机制确保：最频繁使用的嵌入驻留在GPU嵌入缓存中，较为频繁使用的嵌入缓存在CPU内存中，而所有模型参数的完整副本（包括那些很少出现的参数）始终保存在硬盘/SSD上。为了最小化延迟，我们将参数更新以及从更高存储级别（SSD → CPU内存 → GPU内存）迁移缺失参数的过程与密集模型计算重叠。HPS的三级内存架构定义如下：

GPU嵌入缓存（第1级）：这是一个为推荐模型推理设计的动态缓存。它通过巧妙地利用数据局部性，将频繁使用的特征（即热门特征）保留在GPU内存中，从而减少额外/重复的参数移动，以提高嵌入查找性能。GPU缓存支持多种操作符（见第4节），以及动态插入和异步刷新机制（见第6节），以保持较高的缓存命中率。

参数分区（第2级）：在CPU内存中存储嵌入参数的部分副本。它们作为GPU嵌入缓存的扩展，当缓存中不存在所需的嵌入时会被查询。根据应用场景，用户可以选择独立部署或集群部署。在独立部署中，分区可以放置在优化的并行哈希映射（无服务器部署）或本地Redis实例中。分布式部署可以利用多节点Redis配置。每个分区的内容会根据部署中所有推理节点处理的查询异步调整。为了接收在线更新，参数分区可以订阅分布式事件流中的主题。

参数副本（第3级）：为了确保容错性，HPS在每个推理节点的基于磁盘的RocksDB键值存储中保留所有模型参数的完整副本（即模型副本）。如果对相应参数分区的查找请求失败，则会访问此备用存储。因此，如果给定足够的时间预算，HPS部署始终能够为每个查询提供完整的答案。为了保持最新状态，每个节点单独监控分布式事件流，并以其自己的节奏应用在线更新。

4 推理GPU嵌入缓存

在处理在线推理工作负载时，通常无法预知接下来需要哪些嵌入表子集。因此，我们的GPU嵌入缓存被设计为一个通用的动态缓存，它可以通过淘汰旧嵌入来接受新嵌入。

4.1 缓存数据模型

GPU嵌入缓存由如图4所示的三级分层结构组成：槽（slots）、板（slabs）和板集（slabsets）。

图4

• 槽（Slots）：槽是GPU嵌入缓存的基本存储单元。每个槽包含一个嵌入键、相关的嵌入向量以及一个访问计数器。
• 板（Slabs）：现代GPU架构以warp（32个线程的组；[28]）为单位管理和执行代码。通过编写warp感知的程序可以实现峰值性能。因此，我们将32个槽分组为一个板，以便每个warp线程被分配到一个独立的槽。在搜索匹配的嵌入键时，我们使用warp线性探测板。为了确定键是否在板中找到以及找到的位置，我们执行寄存器级的warp内通信（如shuffle、ballot等），以消除分支和内存分歧。
• 板集（Slabsets）：类似于N路组相联缓存中的缓存行被分组为缓存集，板被打包成板集。为了充分利用GPU的大规模并行计算能力，每个嵌入键首先被映射到一个特定的板集，但随后可以占据该板集中的任何槽。这样，线性探测被限制在单个板集内，而不会与独立的板集发生冲突。较小的板集大小可以减少键搜索延迟，但也会导致冲突未命中增加。找到最佳的板集大小以平衡这两个因素非常重要。我们根据经验将板集大小设置为2，适用于当代NVIDIA GPU架构（如Ampere）。

为了最大化GPU资源利用率和推理并发性，推理工作器可以共享同一个嵌入缓存。通过仅允许单个warp对特定缓存操作（如查询和替换）独占访问板集，可以防止竞争条件。这种方法还隐式地确保了线程安全性。由于板集的总数通常远高于每个GPU的最大warp数（数百万对数千），因此互斥不会导致显著的停顿。

4.2 GPU嵌入缓存API

GPU嵌入缓存支持四种API：

• 查询（Query，算法2）：检索嵌入键集合对应的嵌入向量。缺失的键会以列表形式返回，可用于尝试从参数分区中获取这些嵌入。
• 替换（Replace，算法3）：尝试通过先填充空槽来插入嵌入。如果空槽数量不足，则替换最近最少使用（LRU）的嵌入。已存在的嵌入将被忽略。
• 更新（Update，算法4）：首先确定输入键与已缓存键的交集，然后替换相应的嵌入向量。
• 批量导出（Dump batch）：输出当前存储在缓存中的所有嵌入键。

算法2

算法3

算法4

查询、替换和更新共享相同的核心算法（参见算法2、3和4）。对于每个键，分配的处理warp首先使用哈希函数定位包含该键的板集，然后线性探测该板集中的板，以找到匹配的键槽或确定用于插入的空槽/可替换槽（仅适用于替换和更新）。批量导出API较为简单，它只是将当前缓存中的所有键复制到CPU内存中。

所有API都启动异步执行的CUDA内核，即控制流会立即返回到CPU。由于它们在板集级别上是线程安全的（见第4.1节），因此允许并发调用所有API。为了避免频繁启动CUDA内核并提高GPU资源利用率，所有API都接受小批量输入。相应的输入键会公平地分配给warp，并推入warp工作队列中。

4.3 嵌入插入

对于查找失败的情况（即键当前不在GPU嵌入缓存中），会触发缓存插入操作，从CPU内存中的参数分区或本地SSD上的副本中获取缺失的嵌入。如算法1所示，HPS有两种插入模式，GPU嵌入缓存根据当前缓存命中率与用户定义的命中率阈值之间的关系在这两种模式之间切换：

算法1

• 异步插入：如果缓存命中率高于预定义的阈值，则激活异步插入。对于任何缺失的键，立即返回默认的嵌入向量（其值可由用户配置）。实际的嵌入会从更高级别的存储中异步获取并插入到GPU嵌入缓存中，以便在未来的查询中使用。这种惰性插入机制确保了在高命中率的情况下，预测精度损失可以忽略不计。

• 同步插入：同步插入会阻塞管道的其余部分，直到获取到缺失的嵌入。在合理的阈值设置下，同步插入通常仅在预热阶段或模型更新后发生。

5 CPU内存与SSD存储层

为了处理超出GPU内存容量的模型，除了GPU嵌入缓存（第4节）外，HPS在其存储层次结构中还包含了两层额外的存储层。这些层基于系统内存、SSD或网络存储构建，并且高度模块化，以支持各种后端实现。

易失性数据库（VDB）层（图3中的第2层）

VDB层位于易失性内存（如系统内存）中，需要通过NVLink或PCIe总线从GPU访问。与GPU内存相比，系统内存可以以更低的成本扩展。为了进一步扩展，VDB可以利用推理集群中的多个低延迟系统内存。例如，使用我们的RedisClusterBackend VDB模板实现，用户可以将分布式Redis实例用作嵌入的存储后端。因此，VDB实现可以但不限于机器边界。为了分配工作负载，VDB将嵌入表存储组织为分区。分区是嵌入表的非重叠子集，存储在同一物理位置。它们根据共享VDB访问的所有节点处理的推理查询稀疏填充。每个嵌入表的最大大小（=溢出边界）和分区数量是可配置的，并需要权衡。更多较小的分区可以实现更平滑的负载均衡，但每个分区都会增加少量的处理开销。

VDB作为异步缓存运行。如果GPU嵌入缓存报告缺失的键，HPS会接下来查询VDB。与嵌入缓存类似，每个VDB条目包含一个时间戳，指示该条目上次被访问的时间。对于成功检索到的嵌入向量，VDB在返回结果后异步更新此时间戳。缺失的嵌入向量会被调度插入到VDB中，以加速潜在的未来查询。因此，每个嵌入的分区分配是固定的，并由键的XXH64哈希值[4]确定。插入操作是异步进行的，以避免阻塞挂起的查找过程，并随后填充VDB分区。每个分区的逐出策略决定了如果分区超出其溢出边界时应采取的措施。我们实现了多种逐出策略。例如，evict oldest策略会找到并修剪不常访问的键。

持久性数据库（PDB）层（图3中的第3层）

PDB层使用硬盘或SSD永久存储整个嵌入表（即所有模型参数）。因此，PDB有助于提高具有极端长尾分布的数据集的预测精度。PDB层可以作为任意数量模型的备份和最终真实数据源。为了避免键冲突，PDB实现为每个唯一的嵌入表形成单独的键命名空间。

我们的模板实现将嵌入表映射到RocksDB数据库中的列组，存储在每个推理节点的本地SSD上。因此，整个模型数据在每个推理节点中都有副本。通过这种方式，我们实现了最大的容错能力，因为节点故障不会影响其他推理节点完全回答每个查询的能力。即使邻居节点的故障导致附加的Redis VDB宕机，仍可以继续运行。如果没有VDB作为中间缓存，当然可能需要更长的时间才能将缺失键的嵌入向量异步迁移到GPU嵌入缓存中（另见第7节）。然而，假设GPU嵌入缓存能够保持足够高的命中率，客户端应该只会看到推理性能的微小偏差。

6 在线模型更新

到目前为止，我们已经描述了HPS如何组织资源以实现预训练模型的推理。在图5中，我们用红色（→）突出显示了数据流图的这一部分。然而，在许多场景中，推荐依赖于最新信息（例如，社交网络中的用户交互）。在完成训练周期后，增量更新必须传播到所有推理节点以改进推荐。我们的HPS通过专用的在线更新机制实现了这一功能。

图5

易失性与持久性数据库更新

模型训练是资源密集型的，因此由一组与推理集群不同的节点执行。HugeCTR模型的训练集被分割为文件，以最大化嵌入缓存中的局部性。模型通过依次将这些文件加载到缓存中并处理训练片段来进行训练。我们的在线更新机制围绕HugeCTR模型训练构建。它被设计为一个辅助进程（图5中的蓝色[→]数据流图），可以在任何时间点开启或关闭。

一旦训练取得进展，训练节点会将其更新转储到基于Apache Kafka的消息缓冲区[36]。这是通过我们的Message Producer API完成的，该API负责序列化、批处理以及将更新组织到每个嵌入表的独立消息队列中。加载了受影响模型的推理节点可以使用相应的Message Source API来发现并订阅这些消息队列。如图5所示，可以为不同的VDB分区创建单独的订阅。这使得共享VDB的节点也可以在它们之间分担更新工作负载。如果某个节点无响应，其当前分配的任务会转移到其他节点。

应用在线更新不可避免地会增加开销。因此，我们允许每个节点通过后台进程以惰性方式消费更新。更新进程的执行与其他I/O请求对齐。为了控制和调整对在线推理的影响，用户可以限制更新的摄取速度和频率。通过消息缓冲区订阅，更新保证是有序且完整的。因此，在完全处理所有挂起的消息（同步）后，各个数据库级别保证是一致的（即我们保证最终一致性）。我们以惰性方式应用更新意味着在模型更新期间可能会出现轻微的不一致性。然而，在实践中这并不重要，因为模型重新训练的学习率通常非常小。只要优化过程相对平稳，预测性能不应显著下降[2, 20]。请注意，同样的假设也支撑了GPU嵌入缓存查询API的工作原理，如果满足命中率标准，它会为缺失的键返回默认的嵌入值（见第4节）。然而，由于摄取更新不需要停机，因此可以实现持续的模型改进，这使得HPS特别适合与高度活跃的数据源一起使用。

异步GPU嵌入缓存刷新

当推理请求到达时，GPU嵌入缓存需要随时可用。从消息缓冲区到GPU嵌入缓存的持续小更新流会创建难以预测的GPU负载峰值，可能会降低响应时间。因此，我们允许GPU嵌入缓存定期轮询VDB/PDB以获取更新，并在必要时替换嵌入，而不是直接从Kafka摄取更新。刷新周期可配置，以最好地适应训练计划。在使用在线训练时，GPU嵌入缓存可以配置为定期（分钟、小时等）刷新其内容。在使用离线训练时，刷新通过Triton模型管理API[9]发送的信号触发。图3展示了模型更新在GPU嵌入缓存中生效的整个流程：

1. 监控消息流：将更新分发并应用到CPU内存分区（VDB）和SSD（PDB）。
2. 批量导出GPU嵌入缓存键（大小可配置）并将其写入导出键缓冲区。
3. 从CPU内存分区和/或SSD中查找写入导出键缓冲区的嵌入键，
4. 并将相应的嵌入键-向量复制到查询键-向量缓冲区。
5. 将查询键-向量缓冲区下载到GPU设备并刷新GPU嵌入缓存。

7 效果评估

（略）

参考

• 1.https://arxiv.org/pdf/2210.08804

meta在《Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models》提出了DLRM的工程实现。

摘要

深度学习推荐模型（DLRMs）在Meta的许多关键业务服务中得到应用，并且是其数据中心基础设施需求方面最大的AI应用。在本文中，我们介绍了Neo，这是一个软硬件共同设计的系统，用于大规模DLRMs的高性能分布式训练。Neo采用了一种新颖的4D并行策略，结合了表格级（table-wise）、行级（row-wise）、列级（col-wise）和数据并行策略，用于训练DLRMs中的大规模embedding操作（embedding operators）。此外，Neo通过包括混合内核融合、软件管理缓存和质量保持压缩在内的多种关键系统优化，实现了极高的性能和内存效率的embedding计算。最后，Neo与ZionEX配对，ZionEX是一个新的硬件平台，与Neo的4D并行策略共同设计，用于优化大规模DLRM训练的通信。我们在128个GPU上使用16个ZionEX节点的评估表明，Neo在训练已部署生产的12万亿参数DLRM模型方面，性能超过了现有系统高达40倍。

1 引言

深度学习推荐模型（DLRMs）被在线公司广泛使用，包括亚马逊用于其目录中选择商品[35, 37, 58]，Netflix用于展示电影选项[13, 29]，以及谷歌用于显示个性化广告[7, 9, 19]。它们还被标准基准测试组织采用，如MLCommons (MLPerf) [38, 52]。在Meta，我们已经在排序和点击通过率（CTR）预测中广泛使用推荐模型，包括新闻推送和搜索服务[15, 17, 42, 47]。DLRMs是数据中心基础设施需求方面最大的AI应用。

与传统的深度神经网络（DNNs）不同，DLRMs主要包含计算密集型操作（例如，卷积和矩阵乘法），DLRMs结合了计算密集型组件和多达数千个数据密集型嵌入操作符，每个操作符都有不同的资源需求和性能特性[43]。因此，与计算机视觉[8, 18, 59]、自然语言处理[5, 10, 61]和强化学习的同类模型相比，DLRMs通常表现出更低的算术强度和更大的模型大小，实际部署的模型拥有数万亿参数，如图1所示。

图1 在总计算量上比较深度学习模型，以拍拉浮点运算/天(petaflop/s-days)为单位（上部）[45]和模型容量（下部）

现有的针对DNN的软件和硬件解决方案在DLRMs上只能实现次优性能和有限的可扩展性，这是由于以下软件/硬件限制造成的。

在软件方面，现有的深度学习框架通常使用数据、模型或流水线并行化来并行化DNN训练[3, 32, 48]。支持这些策略组合的框架通常为特定的DNN应用而设计[16, 22, 41, 50]。然而，为计算密集型DNN模型设计和优化的现有并行化策略在DLRMs上实现有限的性能和可扩展性。特别是:

• 数据并行化要求每个设备保存整个模型的副本，因此不支持拥有数万亿参数的DLRMs[32]。
• 此外，由于其嵌入操作符的数据依赖行为，DLRM不能直接使用模型并行化或流水线并行化。具体来说，处理不同的训练样本可能需要根据每个样本的分类输入访问不同的嵌入参数。这种数据依赖行为使得在满足所有样本的数据依赖性的同时，静态地将DLRM的可训练参数划分为不相交的子集变得不可行，这是使用模型和流水线并行化的必要条件。

此外，当今的DNN框架旨在优化计算密集型（compute-intensive) DNN计算，忽视了对数据密集型（data-intensive）嵌入操作符的关键优化。具体来说，DLRM包含多达数千个嵌入操作符。这些嵌入操作符的前向处理、反向传播和梯度同步需要在训练迭代中启动数千个CUDA内核，并消耗高达数TB的累积GPU设备内存，引入了显著的运行时开销和内存需求。

在硬件方面，现代硬件平台，如基于GPU的集群，提供了显著的能力提升，但它们并不是为了匹配DLRMs的性能特性而设计的。具体来说，DNN训练的硬件平台通常针对集中式节点间通信（例如，参数服务器[3]）和/或AllReduce通信（例如，Horovod[54]和NCCL[1]）进行优化。然而，如第3节所确定的，高性能和可扩展的DLRM训练需要有效硬件支持多种不同的通信模式，包括AllReduce、AlltoAll、ReduceScatter、OneToMany和ManyToOne。

1.1 我们的方法

我们提出了Neo，这是一个软硬件共同设计的系统，用于快速且可扩展的DLRM训练，它基于三个关键技术构建。

4D并行

为了在DLRM中快速且可扩展地训练大规模嵌入操作符，有效地平衡GPU之间的工作负载分配并最小化通信成本至关重要。我们引入了一种4D并行策略，结合了表格级、行级、列级和数据并行策略，共同优化嵌入操作符的并行性能。此外，Neo还支持在不同级别的硬件层次结构中以递归方式应用4D并行，以进一步提高负载平衡和硬件效率。

高性能嵌入计算

Neo采用了两项新的优化技术，以最小化嵌入操作符的计算成本和内存需求。

• 首先，我们引入了一种混合内核融合技术：将（1）多个嵌入操作符和（2）嵌入计算及其参数更新全部融合在一个CUDA内核中。这是通过共同设计嵌入操作符的优化算法和软件实现来实现的。
• 其次，为了提供足够的内存容量以支持DLRM训练，Neo使用软件管理的缓存机制来利用现代硬件平台的内存层次结构。
• 最后，进一步应用了多种压缩技术[29, 63]来最小化内存需求。

硬件平台设计

我们介绍了ZionEX，这是一个与Neo的4D并行共同设计的新型硬件平台，用于优化分布式DLRM训练的节点间通信。ZionEX通过使用专用的基于融合以太网的RDMA（RoCE）扩展网络，支持集群中所有GPU的全连接拓扑。这种拓扑设计促进了分布式DLRM训练中性能主导的通信工作负载（例如，AlltoAll和ManyToOne）的高性能数据传输。同时，ZionEX支持RDMA和GPUDirect通信协议，并保留了灵活的节点内GPU织物。这使得在ZionEX上能够进行高性能的DLRM训练，同时确保与现有数据中心基础设施的兼容性，允许ZionEX的广泛部署。

结果

我们已经在三个生产环境中部署的不同任务的DLRM上评估了Neo，包括点击通过率预测、排序和参与度，代表了多样化的生产级推荐模型。我们在16个ZionEX节点上的128个A100 GPU上的评估表明，Neo能够处理高达每秒170万次查询，用于训练具有12万亿参数的DLRM，与现有生产中的DLRM训练解决方案相比，速度提升了40倍。消融研究表明，4D并行、高性能嵌入计算和新的ZionEX平台对于实现快速且可扩展的DLRM训练至关重要。

总结来说，我们的贡献是：

• 我们提出了Neo，一个软硬件共同设计的系统，用于快速且可扩展的DLRM训练。Neo在训练具有12万亿参数的大规模DLRM方面，性能超过了现有系统高达40倍。
• 我们提出了4D并行，这是一种结合了表格级、行级、列级和数据并行策略的训练嵌入操作符的方法。
• 我们开发并实现了使用混合内核融合、软件管理缓存和质量保持压缩的高性能嵌入操作符。
• 我们构建了ZionEX，一个与Neo的4D并行共同设计的新型硬件平台，以加速DLRM训练中的多种通信模式。

2 背景

DLRMs通常有两种训练模式——离线和在线，每种模式都有不同的要求。离线训练可以被视为预训练，其中候选模型在足够大的历史数据上进行训练，并期望在部署到当前/未见过的样本时能够泛化。一旦部署，DLRMs继续使用它已经服务过的数据进行在线训练。离线训练受到吞吐量限制，符合更传统的“尽可能快地训练尽可能多的数据”的范式，而在线训练对延迟更敏感，重新训练和更新的频率是一个重要因素。对于在线训练，吞吐量要求较低，因此可能希望使用相对较少的资源。这创造了一个独特的需求，即在能够容忍较低吞吐量的较小规模上训练非常大的模型。

本文专注于对训练吞吐量需求更高的离线训练——每秒处理多达数百万样本（查询），这些样本来自于在合理时间内处理数十PB训练数据。这推动了训练平台的需求，如表1所总结。

嵌入操作符

DLRMs与传统深度神经网络之间的一个主要区别是利用类别型特征（如用户、帖子或页面）。生产中使用的DLRMs通常包含多达数千个类别型特征，每个特征对应一个专用的嵌入操作符。嵌入操作符以一个multi-hot向量作为输入，向量中的每个非零元素触发嵌入表中的完整行检索，其中输入向量的每个索引对应表中的一行。最后，对于给定的输入向量，所有嵌入行通过element-wise pooling组合，如图2所示。

图2 一个embedding operator的Workflow

并行化策略

传统上，用于生产环境中训练DLRMs的是基于参数服务器（PS）的分布式CPU训练系统。具体来说：

• 一方面，MLP模块中的dense参数在训练器之间复制以利用数据并行性。它们的权重使用集中式dense参数服务器通过弹性平均方法SGD进行同步。
• 另一方面，embedding table中的参数被分割并放置在多个PS上以利用模型并行性，因为embedding参数的大小简单地阻止了模型复制。

为了最大化训练吞吐量，使用Hogwild!更新嵌入操作符的参数。此外，读者部署在单独的机器层上，为训练器提供训练批次，如图3所示：

图3 基于分离式参数服务器（Disaggregated parameter-server）的系统

这种基于PS的系统非常适合DLRMs，允许分别扩展不同的组件，并在训练具有不同训练器、参数服务器和读者配置的不同模型时实现平衡的资源利用。此外，系统中的资源在很大程度上是可替代的，这使得数据中心运营成本较低。

然而，需要支持具有数万亿参数的DLRMs，因此大小达到数TB，这对这种方法的可扩展性提出了严重挑战，需要大量增加训练器和参数服务器的数量以满足不断增长的训练需求。这很快变得不可行，由于在大量worker之间增加的异步更新，导致模型准确性因陈旧性而降低。为了解决这些问题，我们构建了一个用于大型DLRMs的高性能同步训练解决方案，将分布式扩展与统计质量解耦。

同步训练系统的高效设计，使我们使用一种新颖的4D并行组合（第4节）用于内存密集型嵌入表，数据并行性用于计算密集型DNN操作符，并在不同组件之间进行流水线处理。这种混合并行性需要AlltoAll通信来处理嵌入查找结果，以及如果输入是从数据库批量流式的，还需要重新分配嵌入表输入，这通常是情况。与用于梯度同步的AllReduce通信不同，AlltoAll通信由于数据依赖性而位于关键路径上，强调了互连和通信原语的性能。此外，DLRMs通常在非常大的数据量上进行训练，这些数据对应于来自各种应用的大多数非结构化和未标记的交互。典型的数据集大小在几个PB的范围内，需要使用常见的分布式网络存储，如Tectonic文件系统。对于训练，这些数据需要被流式传输进来，给主机网络和主机到设备带宽带来额外的压力。

3 概览

图4展示了Neo的概览，这是一个软硬件共同设计的系统，用于快速且可扩展的DLRM训练。本节简要描述了Neo的关键组件。

图4 Neo 概览。图中的每个框表示一个神经网络组件，而框之间的边表示不同组件之间共享的张量。

首先，Neo使用数据并行性来训练计算密集型的DNN层（以橙色显示），并切换到一个4D并行策略，该策略结合了表格级、行级、列级和数据并行性，以高效训练内存密集型的嵌入操作符。

其次，Neo配备了高性能的嵌入操作符实现。这是通过一系列关键的系统优化实现的，包括:

• （1）混合内核融合技术来减少嵌入操作符的计算成本，
• （2）软件管理的缓存机制来利用现代硬件平台的异构内存
• （3）多种质量保持压缩技术来最小化嵌入计算的内存需求

最后，Neo部署在ZionEX上，这是一个与Neo的4D并行共同设计的新型硬件平台，用于优化DLRM训练的节点间通信。

此外，数据I/O是任何训练系统的重要组成部分，特别是随着完全同步训练和加速器的采用。首先，主机到设备的数据传输应该是非阻塞的，并且足够快，不会限制整体训练吞吐量。理想情况下，使用双缓冲或流水线将输入数据传输与训练重叠。其次，尽管将输入数据分布映射到训练器之间的集体通信更快，但这为集体通信的输入和输出数据布局引入了额外的挑战。初步实验表明，这些可能会给关键路径增加显著的延迟。我们将在第7.1节中展示我们如何克服这些实际挑战。

4 4D并行策略

DLRM的一个关键组成部分是嵌入操作符，将在第5节中定义。为了实现嵌入操作符的高性能训练，有效地平衡GPU之间的工作负载分布并最小化通信成本至关重要。我们引入了4D并行策略，它结合了表格级、行级、列级和数据并行策略，共同优化嵌入操作符的并行性能。

图5 具有不同通信成本、负载均衡和内存需求影响的嵌入表分片方案。为了简化说明，此图中省略了底部的MLP（多层感知机）。

表格级并行策略。

最直接的并行方案是：跨GPU分割和并行化多个嵌入表，如图5a所示。表格级并行策略不会进一步分割嵌入表，因此此方案不需要额外处理嵌入表输入索引或聚合的嵌入结果，从而实现最佳的通信效率。然而，表格级并行策略无法处理超过单个GPU内存容量的大型嵌入表，且由于表大小的偏斜，实现的负载平衡通常有限。

行级并行策略。

此方案通过行来并行化大型嵌入表，并将不同的表片段分配给不同的训练器。由于嵌入表输入通过行来索引表，它们需要根据行级并行决策进行分桶并分发到相应的训练器，如图5b所示。此外，多个训练器上的部分结果需要被缩减然后分散到所有训练器以进行下游计算。这要求在前向传递中使用ReduceScatter通信模式。此方案能够很好地处理大型表，并带来更好的负载平衡。然而，通信成本与训练器的数量成线性关系。

列级并行策略。

列级并行策略沿嵌入维度划分嵌入表（见图5c），并将划分后的表视为具有较小嵌入维度的独立操作符。此方案需要为划分后的表复制输入索引。与表格级并行策略相比，它保持了相同的流程和通信模式（AlltoAll）。列级并行策略的一个关键优势是能够实现更细粒度的并行策略，特别是对于大型表。然而，它仅在大型嵌入维度下表现良好，并增加了输入索引的负载，这些索引必须复制到所有节点的列片段。此外，由于列级划分的表的行分布在不同的训练器上，使用这些表的独立行更新引入了额外的参数，每个行片段一个参数，而不是使用稀疏优化器时整个行只有一个单一值（见第5.1节了解更多细节）。

数据并行策略。

DLRM往往有广泛的表大小范围，而表格级、行级和列级并行策略适用于相对较大且无法复制的嵌入表。对于较小的表，数据并行策略能够实现更好的性能，因为数据并行策略在前向传递中不涉及任何通信（见图5d）。因此，对于小型嵌入表，Neo将embedding表视为dense参数并在所有训练器上复制它们。对于数据并行嵌入表的聚合嵌入，不再需要AlltoAll。相反，需要AllReduce来同步所有副本。因此，这取决于聚合嵌入的AlltoAll成本与整个表的AllReduce成本之间的权衡。通常，行数较少的小型嵌入表是数据并行策略的良好候选者。这些表的输入索引作为数据并行输入传递，不再需要重新分配。

4.1 并行化算法

Neo支持在单个嵌入操作符的粒度上应用4D并行策略，以最大化灵活性。实践者可以混合使用上述原语，以确定划分嵌入操作符的最佳策略。此外，Neo还支持在不同级别的硬件层次结构中以递归方式划分嵌入操作符，以进一步提高工作负载平衡和硬件效率。例如，表格级然后行级方案首先将一组表分配给特定节点，然后在该节点内按行划分表。这种层次并行方案通过充分利用快速的GPU互连并减少节点间通信，提高了硬件局部性。

为上述每种并行方案定义了成本函数，可以探索放置算法以最小化worker之间的成本差异。成本函数是通信开销和训练器之间的负载不平衡的组合。通信开销使用消息量作为代表性指标计算，消息量越大对应成本越高。这在很大程度上准确地捕捉了吞吐量成本，对于延迟测量值，作为固定附加成本纳入。我们通过使用每个训练器的嵌入访问大小来估计负载不平衡，这可以近似为每个训练器的嵌入表数量×全局批量大小×每个样本的平均索引数×嵌入维度。这两种成本的组合为我们提供了通信和负载不平衡的合理估计。进一步，我们为每个单独的成本引入了标量权重，可以根据不同的系统规格进行调整，以获得更准确的估计。

我们实现并评估了两种多项式时间启发式算法作为概念验证。第一个是一个简单的贪婪启发式算法，它将可用方案的成本按降序排序，并首先分配最大的片段，每个worker一个。然后，贪婪算法遍历所有剩余的片段，并将最高成本分配给成本总和最小的节点。第二个启发式是最大差分方法（也称为Karmarker-Karp算法[26]）。主要思想是从输入中取出两个最大的数字，并用它们的差替换它们。它直接减少了总和的差异，通常优于贪婪启发式。

4.2 流水线

尽管使用GPU作为主要计算资源在模型评估内提供了有限的流水线机会，我们通过流水线化批间数据移动并与计算重叠通信来提高GPU利用率。

当批次𝑖正在被评估时，相同的GPU可以开始使用单独的流接收和分发批次𝑖 + 1。为了最小化干扰，我们将批次𝑖 + 1的输入AlltoAll与批次𝑖的顶部MLP的前向传播重叠，其中不涉及通信。此外，我们将聚合的嵌入AlltoAll与底部MLP的前向传播重叠，以隐藏延迟。

5 嵌入优化

优化DLRM的嵌入操作符（见第2节）的运行时性能需要解决两个关键挑战。

• 首先，嵌入操作符的前向处理、反向传播和梯度更新需要在每次训练迭代中启动数千个GPU内核，引入了显著的GPU内核启动开销。
• 其次，一些嵌入操作符可能包含高达数十亿的参数，无法适应单个GPU的设备内存。

我们引入了三种新技术来减少嵌入操作符的计算成本和内存需求。

• 首先，我们引入了一种混合内核融合技术，以最小化CUDA内核启动开销，并允许每个GPU工作器只启动两个内核（即一个用于前向传播，一个用于反向传播和参数更新）。
• 其次，对于嵌入操作符的并行计算，我们提出了列级并行策略和行级并行策略，除了数据和模型并行策略之外。这四个并行维度的组合使Neo能够支持高达数万亿参数的嵌入表。
• 最后，Neo利用一系列内存节省技术，利用ZionEX平台的内存层次结构，确保DLRM有足够的内存容量。

5.1 内核融合

Neo使用混合内核融合机制来最小化执行嵌入计算的CUDA内核启动开销。

• 首先，与为每个嵌入表应用单独的嵌入查找不同，Neo将同一GPU上的多个嵌入查找融合到单个CUDA内核中（图6a），这提高了并行性和带宽利用率，并减少了在GPU上启动多个CUDA内核的开销。
• 其次，Neo还将反向传播与稀疏优化器融合，以进一步减少内核启动开销，并避免将梯度具体化到嵌入表中。这种融合的关键挑战是避免来自不同训练样本的梯度更新之间的潜在竞态条件，以及处理诸如AdaGrad[11]、LAMB[66]和Adam[27]等高级优化器中的非线性。例如，图2中的样本1和2都有助于嵌入向量1和6的梯度。如果不聚合直接将这些梯度发送到非线性稀疏优化器，将导致嵌入表的错误更新。为了保证正确性的同时最大化性能，Neo通过行对梯度进行排序，以便对同一嵌入行的梯度由单个CUDA线程块处理，如图6b所示。
• 随后在每个CUDA线程块内使用更快但更小的GPU共享内存进行梯度聚合。

图6

Neo的嵌入操作符混合融合技术带来了三个性能优势。

• 首先，Neo通过避免为嵌入梯度分配GPU设备内存来减少嵌入操作符的内存需求。
• 其次，通过使用GPU共享内存来保存中间嵌入梯度，最小化了对GPU设备内存的内存访问。
• 最后，内核融合通过高达7倍的性能提升，改善了嵌入计算的整体性能，与原生实现相比。

优化的嵌入操作符实现作为FBGEMM库∗的一部分开源，并与PyTorch集成。

5.2 管理内存层次结构

对于具有高达数万亿参数的DLRM，嵌入表太大，无法完全适应单个GPU。我们利用ZionEX平台的多个内存层次结构，包括HBM、DRAM和SSDs，以及扩展到多个节点以增加聚合容量，确保模型有足够的内存，更快的内存作为后续层的软件缓存。Neo的层次内存管理策略特别适用于DLRM的在线训练，由于吞吐量要求较低，因此可以使用较少的节点来训练原始的大型模型，如第2节所述。

管理内存层次结构的一种方法是CUDA的统一内存（UVM）[44]，它为不同类型的内存提供单一的内存地址空间，并自动替换和逐出未使用的页面。然而，嵌入操作符中的随机表查找需要在单个嵌入行的粒度上缓存和替换未使用的参数，这使得直接使用UVM对于DLRM来说是不够的。需要额外处理查找以确保性能不受频繁的主机到设备传输的限制。相反，Neo使用定制的32路集合关联软件缓存[64]，使用最近最少使用（LRU）或最少频繁使用（LFU）缓存替换策略，其中关联性与GPU的warp大小相匹配。这使得可以对缓存和替换进行细粒度控制，允许针对目标模型特性进行调整。请注意，UVM受PCIe带宽限制，而Neo的软件缓存可以弥补PCIe和HBM之间的带宽差距（50倍差异）。与UVM相比，软件缓存将DLRM工作负载的端到端性能提高了约15%。

为了进一步减少嵌入操作符的内存需求，Neo还采用了先前工作中引入的多种压缩技术，如逐行稀疏优化器[14, 62]、使用高精度缓存支持的低/混合精度训练[63]和高级分解技术[29]。

逐行稀疏AdaGrad首次在[14]中引入，然后在[62]中进一步阐述。在逐行稀疏AdaGrad中，每个元素的时刻估计应用于整个嵌入行。对于每一行，它是一个单一的缩放因子，通过添加行中梯度的平均平方和来更新。通过这种方式，我们将动量保持为一个1D张量，有H个元素，而不是H×D的2D张量，其中H和D分别是嵌入表中的行数和每行的元素数。

6 ZIONEX: 硬件平台设计

我们首先在第6.1节描述了我们之前用于DLRM的硬件平台的局限性。第6.2节介绍了ZionEX，一个为DLRM设计的新型硬件平台。我们还概述了在ZionEX开发中使用的设计原则。

6.1 之前的平台：Zion

2019年推出的Zion是我们之前针对DLRM训练的高性能硬件平台的工作。尽管Zion在单节点级别提供了显著改进的能力，但作为一个分布式平台，它未能扩展以满足迅速增长的DLRM训练需求。我们批判性地评估了它的局限性，但其他基于类似设计的平台也存在相同的局限性；我们在第9节讨论了这些平台。

图7a显示了一个Zion节点的架构，它有8个CPU插槽，1.5TB内存，8个GPU和8个网络接口卡（NIC）。它通过（1）将DLRM的计算密集层（例如，MLP）卸载到GPU上，以及（2）利用CPU处理大型嵌入操作符在相对便宜的DRAM上，而不是HBM，以在单个节点上容纳TB级DLRM，提供了强大的异构超级节点设计。

图7

然而，这种异构设计引入了许多软件设计和性能挑战。例如，平衡CPU和GPU上的工作负载以确保最大重叠至关重要。这需要在CPU和GPU之间进行精细的流水线处理，并使用精确的成本模型将DLRM划分为细粒度任务。此外，DLRM的异构训练还引入了非平凡的运行时开销，例如增加了CPU和GPU之间的数据传输和节点间通信。最后，Zion的一个关键缺失组件是每个NIC直接连接到一个CPU。因此，所有节点间通信（例如，梯度同步和张量转换）都需要CPU干预和额外的GPU-CPU传输。此外，这些NIC连接到常见的数据中心网络基础设施，引入了网络拥塞的开销和干扰，并且受限于使用更数据中心友好的拓扑结构和协议（TCP/IP），这对于分布式训练是次优的。尽管每个Zion节点都配备了8个100Gbps NIC带宽，但实际上我们发现由于网络开销，很难扩展到多个节点。随着DLRM模型大小需求的增加，Zion无法很好地扩展并充分利用强大的硬件资源。

6.2 ZionEX

为了解决这些不足，我们引入了ZionEX，我们已经设计它比之前的Zion平台更具可扩展性，并提高了网络能力，同时保留了其灵活性和核心优势，例如OAM外形因素、模块化设计和灵活的节点内加速器结构。通过所有这些改进，ZionEX在支持增加模型复杂性和提高训练性能方面带来了数个数量级更高的能力。这最好通过比较每个平台支持的最大模型复杂性（以FLOPS/样本计）和实现的训练吞吐量来说明，这可以被视为标准化的有效性能。对于ZionEX，实现了1.2 MQPS的吞吐量，模型复杂性为638 MFLOPS/样本（见表3），这转化为766 TFLOPS/s的有效性能，还有额外的余地上升到数PETAFLOPS/s。而Zion在ZionEX上支持的最大模型复杂性不到一半（约250 MFLOPS/样本），吞吐量更低（约0.25 MQPS），因此最大可实现的有效性能降低了10倍以上，仅为63 TFLOPS/s。图7b显示了整体系统架构。我们简要强调了ZionEX的核心技术原则：

可扩展性。Zion和ZionEX都支持DLRM的异构训练，但最显著的区别是ZionEX设计了足够的扩展和扩展网络能力。如图7b所示，ZionEX为每个通过PCIe交换机连接的GPU配备了专用的RDMA over Converged Ethernet (RoCE) NIC，以允许专用的节点间连接（与常见数据中心网络隔离），并重要地支持更高效的RDMA/GPUDirect通信协议。这些ZionEX节点可以通过专用后端网络连接，形成一个分布式可扩展训练的集群。ZionEX的可扩展设计允许扩展后端网络，连接数千个节点，形成一个数据中心规模的AI训练集群。

高性能。作为一个扩展解决方案，我们将整个DLRM卸载到GPU上，充分利用大规模并行性和高内存带宽来加速MLP和嵌入计算。为了传输张量和同步梯度，每个GPU都可以直接通过专用的低延迟高带宽RoCE NIC与不同节点上的GPU通信，不涉及主机CPU。此外，ZionEX还有一个前端NIC连接到每个CPU。数据摄取通过常规前端网络和PCIe进行，不干扰激活或梯度。主机CPU仅用于设置输入批次和组织训练过程。

能力。通过ZionEX，我们确保平台与现有基础设施兼容，并可以在我们的数据中心内广泛部署，不会造成重大中断。这对于能够有效利用平台的能力并使其随时可用于各种应用和用例至关重要。我们通过使ZionEX平台符合标准的Open Rack规范来实现这一点，这涵盖了与其他基础设施组件的兼容性，如电源、冷却、机械和布线。此外，设计平台为模块化，并依赖基于开放标准技术，例如基于以太网的网络结构，用于高性能扩展解决方案。

图7c显示了整体训练平台，以及分离的数据摄取服务。这支持从网络存储（如Tectonic）流式传输输入数据，并以分布式方式执行轻量级数据预处理操作。以便数据摄取不是端到端训练的瓶颈，并确保在向ZionEX训练器提供数据时有足够的吞吐量。

7 实现

我们详细描述了上述用于DLRM的高性能可扩展训练的实现。我们使用PyTorch [48]构建了一个高性能训练软件栈，通过ATen库为大多数深度学习操作提供高效的CUDA实现，并通过PyTorch DistributedDataParallel库自动处理参数复制和梯度同步，以及通过重叠反向传播和AllReduce实现。我们已经启用了以下组件以实现高效的DLRM训练。

7.1 数据摄取

数据摄取是确保端到端训练性能的关键组件，特别是对于DLRM，它们通常处理的数据量比其他典型的DNN模型大得多。我们观察到，如果未经优化，数据摄取可能会引入显著的延迟，并为流水线带来非平凡的开销。

最初为分布式异步CPU设置设计的我们的读取器和数据预处理模块将每个稀疏特征的偏移量和索引存储在单独的张量中，每个嵌入表一个。因此，具有数百个嵌入表的DLRM可以轻松地在每次迭代中获得数千个输入张量，这转化为从CPU ↔ GPU传输的重大开销，并且是之前Zion平台的主要瓶颈之一，如第2节所述。

为了克服这一实际挑战，我们共同设计了数据预处理模块，使用组合格式，其中使用长度而不是偏移量，并将不同嵌入表的输入简单地连接起来。使用组合格式的好处是两方面的：（1）它通过合并小传输来优化CPU-GPU传输；（2）它可以直接被嵌入内核消耗，无需额外的布局转换。我们进一步通过使用固定内存来优化输入数据传输，以避免额外的复制。

有了组合格式，我们开发了一个模块，根据分片策略高效地分发嵌入表输入。在表格级分片（如图5a所示）的情况下，需要一个AlltoAll来将全局批次分发给每个工作器的本地表。由于索引的大小取决于长度的值，通信实际上是先进行长度的AlltoAll，然后进行索引的AlltoAll。在有𝑊个工作器，𝑇个本地表和𝐵个本地批次大小的设置中，这给我们提供了（𝑊，𝑇，𝐵）顺序的索引，需要进一步排列为（𝑇，𝑊，𝐵）以供嵌入内核消耗。我们已经开发了自定义的GPU内核，用于排列、分桶和复制，以实现表格级、行级和列级分片方案的最大吞吐量。模型检查点也面临类似的挑战，需要足够频繁地能够写出更大的模型，同时不成为训练的开销，如这篇最近的论文[12]所概述的。

7.2 通信原语

高性能的集体通信是DLRM训练表现良好和可扩展性的关键。PyTorch提供了进程组（PG）接口，用于集体操作——一个抽象的平台/集体库不敏感的API。DLRM直接（对于Alltoall）或间接（通过DDP对于Allreduce）使用这个API[32]。我们使用NVIDIA的集体通信库（NCCL）作为我们的主要集体通信库，因为它有效地使用RDMA和NVLINK以获得最佳性能。我们将PyTorch NCCL进程组实现扩展到支持使用NCCL Send/Recv原语的Alltoall/Alltoallv集体操作（需要NCCL 2.7.3或更高版本）。

8 评估

我们提供了生产模型端到端训练的结果，以及操作级别的性能分解。

8.1 实验设置

表2总结了配备8个NVIDIA A100 GPUs的单个ZionEX节点的聚合能力。节点中的8个GPU提供了总共320GB的HBM，聚合内存带宽为12.4TB/s。4个插槽的CPU提供了1.5TB的内存和320GB/s的带宽。在网络能力方面，GPU通过高带宽NVLink进行节点内GPU通信，每个GPU都有一个专用的200Gbps RoCE NIC用于节点间通信。我们在实验中使用了16个ZionEX节点的集群，总HBM容量为5TB。

8.2 端到端训练

我们报告了三个在生产中部署的DLRM的结果，这些DLRM用于不同的任务，包括点击通过率（CTR）预测、排序和参与度。表3列出了这些候选模型的高级特性。模型A代表了大型和复杂的DLRM，它们强调了Neo的计算能力和通信带宽，每个样本使用显著更高的FLOPS和大量的嵌入。模型F提出了一个不同的实际挑战，尽管每个样本的FLOPS很低，嵌入表的数量很少，但它有一个巨大的单一表，无法适应单个GPU的设备内存。最后，模型I代表了中等规模的DLRM，它们通过高平均嵌入池大小强调内存带宽。这些目标模型在集群中最多在16个ZionEX节点（128个GPU）上进行训练。模型质量以归一化熵[20]评估，训练吞吐量以每秒查询数（QPS）衡量。

图8

首先，我们使用模型A来演示训练质量，因为它也可以在分布式CPU平台上训练。如图8所示，尽管使用了更大的批量大小（64K vs. ~150），在ZionEX上同步大批量训练提供了相当或更好的模型质量（两者都使用调整过的超参数）。在相同的配置下，Neo在16个节点上的128个GPU上实现了1.2 MQPS，与我们之前的一代分布式CPU异步训练平台相比，速度提升了40倍，后者使用了45个参数服务器和15个训练器。以前的解决方案在不损害训练质量的情况下无法进一步扩展，而在ZionEX上完全同步训练允许在16个节点之外进行扩展，甚至可以使用更大的批量大小。

8.3 扩展性能

图9显示了在保持每个GPU批量大小不变的情况下，使用多达16个节点的模型A和模型I的归一化训练吞吐量。虽然数据并行训练的工作负载随着扩展保持不变，但由于模型并行策略，每个GPU的嵌入表数量随着扩展而减少。出于同样的原因，每个GPU为其本地表处理整个全局小批量，这与扩展成比例增加，并补偿了减少的表，使得这仍然是一个弱扩展实验。要在较小的节点数量上运行，我们减少了嵌入表的基数，并散列输入以适应减少的行数。这个缩小版本的模型有效地减少了模型大小，对性能特性的影响最小/没有影响，因此用于研究扩展性能。

图9

从图中可以看到，在较大的节点数量上，模型A的扩展效率约为50%，模型I约为75%。尽管模型A和模型I在考虑目标本地批量大小时在有效FLOPS和内存需求方面非常接近，但模型A有更大的完全暴露的AlltoAll延迟。这是因为更多的嵌入表增加了AlltoAll负载，并且混合维度使得同时平衡嵌入计算和AlltoAll通信更加困难。因此，模型A在扩展时受到AlltoAll效率降低的影响更大。

为了更好地理解扩展性能，我们在图10中提供了模型A的序列化和暴露训练迭代延迟的分解。比较序列化和暴露延迟，CPU到GPU传输（即HtoD）完全隐藏，暴露的通信延迟远小于序列化AlltoAll和AllReduce延迟的总和。这证明了Neo的流水线优化有效性，可以重叠通信与计算（见第4.2节）。

图10

随着节点数量的增加，我们观察到AlltoAll和AllReduce延迟增加。由于大多数AlltoAll通信在关键路径上，增加的AlltoAll成本直接影响暴露的通信和整体训练延迟。虽然AllReduce在多达16个节点上大部分被隐藏，但增加的AllReduce延迟和不变的计算延迟表明，一旦后向传递中的松弛完全被更高节点数量和/或更快的计算用完，AllReduce可能成为瓶颈。

8.4 训练吞吐量优化

以模型A作为案例研究，我们详细说明了各种优化及其在实现高达1.5 MQPS（如图11所示）中的贡献。此外，我们使用附录B中描述的性能屋顶线建模方法来建立可实现性能的上限，并确认报告的吞吐量在理论估计的15%以内。模型A在128个GPU上的基线性能低于700 KQPS。进一步的分析揭示了不同GPU之间嵌入查找延迟的巨大差异，表明存在严重的负载不平衡。通过结合表格级、列级和数据并行策略来处理约1000个嵌入表的≈1000s，将它们分配到128个GPU上，从而缓解了这个问题。请注意，尽管列级并行策略引入了额外的输入AlltoAll成本，但更好的负载平衡的好处超过了开销，总体QPS提高了20%。然而，扩展效率仍比理想的线性扩展低约30%。

图11

如前所述，限制扩展效率的两个主要问题是：（1）负载不平衡和（2）增加的AlltoAll延迟。对于模型A，仅使用HBM进一步平衡负载特别具有挑战性，因为模型大小在TF32下接近128个GPU上的5TB聚合HBM容量。在扣除PyTorch框架和NCCL在每个等级上保留的内存后，Neo几乎没有空间探索放置策略。为了缓解这个问题，我们使用较低精度（FP16）的嵌入表，将模型大小减少了2倍。虽然这本身并不直接提供吞吐量好处，但Neo可以利用这个空间更好地平衡。结果，由于改善的负载平衡，训练吞吐量又增加了20%。

接下来，为了解决增加的AlltoAll延迟问题，我们采用了[65]中提出的量化集体通信，这直接减少了通信量。对于模型A，我们验证了在前向AlltoAll中使用FP16和在后向AlltoAll中使用BF16几乎提供了30%的速度提升，而没有训练质量损失。最后，我们将全局批量大小从64K增加到256K。这直接增加了激活大小，有助于更好地饱和GPU和通信带宽，同时与其他所有优化相辅相成。在适当调整优化器/超参数后，我们能够实现与训练质量相当的训练，但需要更全面的实验，因为DLRM的大批量训练研究得不够充分，将成为未来工作的一部分。总的来说，这些技术相比使用64K全局批量大小的TF32训练，训练吞吐量提高了87%。

8.5 模型容量限制研究

我们以模型F为例，在原型系统上推动模型容量。与模型A或模型I不同，有效训练模型F提出了2个不同的挑战。首先，模型F有12T参数，使用简单的训练方法，模型F很容易需要高达96TB的内存，远远超过了16个节点集群上的总内存。其次，模型只有几个巨大的嵌入表，每个表有约100B行和256列，每个表需要多节点的GPU和主机内存来训练。

为了将模型适配到16个节点上，我们首先应用逐行稀疏AdaGrad优化器到嵌入表，这将优化器状态从每个元素减少到每个嵌入行。然后我们在嵌入表上使用FP16精度[67]。这两个优化共同将模型内存占用从96TB降低到24TB，刚好适合4TB HBM + 24TB DRAM内存层次结构。在巨大的嵌入表上，我们启用逐行分片将表分布到多个节点，并调整训练流程使用AlltoAll与桶化和ReduceScatter，如图5b所示。在启用UVM并使用HBM作为缓存的情况下，我们能够以高达1.7 MQPS的吞吐量训练模型F，展示了我们HW/SW共同设计解决方案推动超越当前最先进技术的能力。

9 相关工作

研究人员提出了各种系统级创新来应对极大模型带来的挑战。

• DeepSpeed [50]在所有节点上完全分割模型参数、梯度和优化器状态，并使用检查点分区和重新物化[21, 28]来动态重建必要的状态，从而大幅减少内存使用。
• GShard [31]通过在张量级别标注并行策略，训练一个巨大的翻译模型，该模型跨加速器进行分割。
• FlexFlow [22]使用自动搜索来发现图中最佳的操作符并行策略。在自动并行化这一方向上，这些最近的论文[39, 60]使用最优合成和强化学习来找到优化的设备放置，以进一步提高并行性，无需手动干预。

然而，上述这些通用系统并非专门为高度稀疏的推荐模型设计。为此：

• 阿里巴巴引入了XDL [23]，这是一个为高维稀疏数据设计的工业级训练系统。XDL包含了诸如层次样本压缩、工作流流水线、零拷贝和CPU绑定等优化，以提高模型稀疏部分的训练效率。
• Kraken [62]针对更高效的在线训练，通过解耦键值获取和嵌入、与ML领域知识共同设计的缓存逐出策略、针对模型的稀疏和密集部分的内存高效优化器，以及允许推理服务器和参数服务器独立扩展的非共址部署模型。
• [25]通过无锁嵌入表更新、调整循环平铺来优化基于CPU的DLRM训练，AlltoAll通信原语，以及利用FP32和BFloat16中的位别名优势来减少内存占用的新split-SGD实现。
• 百度的AIBox [70]采取了不同的方法进行水平扩展，专注于在单个节点上适应大型推荐模型的训练。AIBox通过流水线网络、磁盘和CPU/GPU任务隐藏服务延迟，减少模型更新开销，并通过分组哈希方案和多级内存哈希系统提高SSD寿命。

由于通信性能已成为集群和数据中心规模分布式训练的主要瓶颈，因此对通信性能的关注越来越多。

• BytePS和ByteScheduler [24, 49]利用空闲CPU和网络资源以及更好的通信调度来提高参数交换效率。然而，在每个作业跨越多个节点的同质训练集群中，寻找和利用空闲网络资源的机会减少，导致这种方法的次优使用。
• SwitchML和ATP [30, 53]利用可编程网络交换机在数据中心环境中执行网络内聚合，以减少跨机架带宽。
• [6, 36]发现并利用数据中心网络的局部性，并通过学习和最优合成形成优化和动态的聚合路由。
• 这些论文[33, 34]通过使用各种量化方案来减少通信量，以解决通信开销问题。

参考

• https://arxiv.org/pdf/2104.05158