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(𝜎)在它的全局最小值上。

#

介绍

字节在《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 PStraining 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采取了类似的方法来弹性地接纳和逐出特征,同时它有一个更轻量级的参数同步机制来保证模型质量。

参考

介绍

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 效果评估

(略)

参考

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]通过使用各种量化方案来减少通信量,以解决通信开销问题。

参考

facebook在2020年《Understanding Capacity-Driven Scale-Out Neural Recommendation Inference》提供了它们的可扩展inference方案。看下它的实现:

摘要

深度学习推荐模型已经发展到TB级别传统的服务方案——将整个模型加载到单个服务器上——无法支持这种规模。支持这种规模的一种方法是采用分布式服务,即分布式推理,它将单个大模型的内存需求分散到多个服务器上。

这项工作为系统研究社区开发新型模型服务解决方案迈出了第一步,鉴于巨大的系统设计空间。大规模深度推荐系统是一种新颖的工作负载,并且至关重要,因为它们在数据中心中消耗了高达79%的推理周期。为此,这项工作描述并刻画了使用数据中心服务基础设施进行扩展的深度学习推荐推理。这项工作特别探讨了延迟受限的推理系统,与最近其他工作中以吞吐量为导向的训练系统进行了比较。我们发现,分布式推理的延迟和计算开销,主要是由于模型的静态嵌入表分布输入推理请求的稀疏性造成的。我们进一步评估了三个类似DLRM模型的嵌入表映射策略,并详细说明了在端到端延迟、计算开销和资源效率方面具有挑战性的设计权衡。总体而言,我们观察到在数据中心规模的推荐模型以分布式推理方式提供服务时,延迟开销仅为边际——在最佳配置情况下,P99延迟仅增加了1%。延迟开销主要是由于所使用的商用基础设施和嵌入表的稀疏性造成的。更令人鼓舞的是,我们还展示了分布式推理如何能够提高数据中心规模推荐服务的效率。

I. 引言

深度学习推荐模型在提供高质量互联网服务方面发挥着重要作用,并且最近开始受到系统社区的关注 [1]–[8]。推荐工作负载已被证明在2019年占Facebook数据中心所有AI推理周期的79% [1],并被纳入MLPerf基准测试工作 [9]–[11]。这一数据中心工作负载的重要性值得对其性能、吞吐量和可扩展性给予更多关注。这类独特的深度学习工作负载面临的一个主要挑战是其规模的不断增长。图1展示了Facebook推荐模型的增长率。此外,百度(Baidu)和谷歌(Google)也部署了规模为1-10TB的模型 [12]–[14]。这种增长是由对更高准确性的需求推动的。众所周知,增加模型参数数量可以提高泛化能力和准确性,并且在更大的训练数据集下是实现最佳结果的必要条件 [15]–[20]。这一现象在多种深度学习应用中普遍存在,包括具有大嵌入表的语言模型——类似于深度推荐系统。本文不探讨增加模型规模对准确性的影响,而是专注于支持生产环境中大规模深度推荐模型的系统和性能影响。

图片名称

图1

深度推荐模型的规模(以及内存容量)主要由嵌入表(embedding tables)决定,嵌入表表示学习到的稀疏参数。每个稀疏输入通过哈希函数映射到其对应嵌入表中的一个或多个索引,每个索引映射到一个向量。索引到的嵌入向量随后通过池化操作(如向量求和)进行组合。增加哈希桶大小和嵌入表数量可以提高嵌入捕获的信息量,从而直接提升推荐模型的准确性,但这也会导致内存占用的相应增加 [21]。减少哈希桶大小可以限制模型规模,但必须谨慎操作以保持模型准确性。

随着大型推荐模型的内存需求超过单个服务器的内存容量,需要解决方案来限制模型规模或增加有效内存容量。模型压缩技术可以用于限制模型规模,但可能会降低模型准确性 [22]。增加DRAM容量也是一种方法,但在超过个位数TB范围后无法扩展。即使如此,支持大容量内存也需要更大的服务器机箱。另一种解决方案是从更高容量的存储设备(如固态硬盘SSD)中按需分页加载模型,但这需要快速的SSD来满足延迟要求。

在数据中心中,额外的硬件需求是不可取的,因为这会增加管理异构资源的复杂性。例如,具有特殊配置的集群在高活动期间无法轻松扩展资源,或在低活动期间无法高效缩减资源。这对于受昼夜流量模式影响的工作负载尤其明显 [23]。与上述方法相比,分布式服务范式可以在现有基础设施上实现,通过将模型拆分为独立子网络并在通用CPU平台上运行。当然,这种方法并非没有权衡。额外的网络延迟、网络负载以及对额外计算节点的需求都会随之而来。然而,在大型同构数据中心中,与定制化、非常规硬件平台相比,部署和扩展这些资源更为容易。

因此,我们为研究社区提供了深度学习推荐系统分布式推理的详细描述和特征分析。这是应对模型规模增长挑战的首个可扩展解决方案,并为进一步优化提供了基线系统。值得注意的是,它与最近以吞吐量为导向的训练系统不同,后者没有推理服务中的延迟限制 [12], [15]。该实现基于真实输入和缩放的深度学习推荐模型进行了特征分析。在此服务范式中探索了三种模型并行化策略。这一具有挑战性的新颖工作负载为系统研究社区在数据中心规模的机器学习领域提供了新的机会。

本文的贡献如下:

  • 据我们所知,这是首个描述大规模神经推荐推理分布式服务基础设施的工作。
  • 我们深入分析了分布式推理对端到端推荐服务延迟、尾部延迟和算子计算开销的影响。我们进一步研究了通过模型分片(sharding)进行嵌入表放置的影响,并将研究结果置于数据中心服务环境的背景下。
  • 最后,我们设计了一个跨层的分布式检测框架,用于性能调试和优化分析,以量化远程过程调用(RPC)服务和机器学习框架的性能开销。

本文的结构如下:第二部分简要介绍深度学习推荐系统。第三部分描述分布式推理的应用和实现,特别关注用于生成分布式模型的分片策略。第四部分描述了用于收集工作负载测量的自定义跟踪框架。第五部分提供了我们特征分析中使用的模型、平台和输入的详细信息。第六和第七部分展示了我们的特征分析结果,并将其置于数据中心服务环境的背景下。我们为系统设计者提供了指导和建议。第八部分讨论了相关工作。最后,第九和第十部分总结了我们的结论。

II. 大规模推荐推理

推荐任务是从一组产品或内容中为用户推荐或建议可能感兴趣的项目。例如,服务可以根据用户明确喜欢的内容和隐式消费的内容推荐新的视频片段、电影或音乐。模型的准确性是衡量用户对推荐结果的兴趣和满意度的抽象指标。传统上,基于邻域的技术(如矩阵分解)通过提供基于与其他用户的相似性或偏好项目的相似性的推荐,取得了良好效果 [24], [25]。然而,最近的推荐系统使用深度神经网络将各种稠密(dense)和稀疏(sparse)输入特征组合成一个更通用的预测模型 [1], [3], [26]–[29]。稠密特征的一个例子是用户的年龄;稀疏或分类特征的一个例子是用户喜欢的网页,模型的输出是候选项目输入的排名。图2a展示了这种深度学习推荐模型架构的简化概述。

图片名称

图2

目前,推荐推理主要在CPU上执行,与其他深度学习工作负载中流行的异构系统不同 [1], [2]。这是因为,与其他推理任务中使用的GPU和其他AI加速器相比,推荐模型中的(1)稀疏性,(2)推荐算法的不断演变,以及(3)延迟限制的吞吐量约束,使得推荐推理在大规模上难以高效地在AI加速器上运行。吞吐量(每秒查询数,QPS)是推理的关键目标,但延迟约束同样重要。为了提供令人满意的用户体验,推荐结果需要在规定的时间窗口内返回。这种严格的延迟约束定义了服务级别协议(SLA)[1]。如果无法满足SLA目标,推理请求将被丢弃,转而使用可能质量较低的推荐结果,这可能会恶化用户体验。为了在相应的SLA约束内最大化吞吐量,采用了各种技术,例如批量大小调整和模型超参数调优。

1) 稀疏特征操作

稠密特征通过全连接(FC)层进行处理,而稀疏特征(如图2a所示)在进一步特征交互之前会经过一次转换。稀疏输入被转换为访问ID列表或哈希索引,这些索引指向嵌入表。嵌入表的大小(桶的数量和向量维度)是一个可调的超参数。通常,每个稀疏特征有一个嵌入表,但为了节省内存资源,特征可以共享表。例如,在一个N × M的表中,嵌入维度为M,具有K个索引的稀疏特征输入将生成K个长度为M的嵌入向量。池化操作(如求和或拼接)将沿第一维度折叠矩阵,生成1 × M或1 × (MK)的结果,用于特征交互。在本工作使用的Caffe2框架中,嵌入表操作符称为SparseLengthsSum(SLS)。通常也将相关操作符家族称为SLS。由于可能的输入数量巨大,嵌入表的大小受到限制,可能会发生哈希冲突。例如,一个维度为32的FP32嵌入表,用于30亿唯一用户,将消耗超过347GB的内存。为了在单个服务器上可行地使用,需要减小表的大小。

2) 模型规模的显著增长与大嵌入表

为了提高模型准确性并利用稀疏特征中丰富的潜在信息,深度学习推荐模型的规模迅速增长,如图1和最近的研究所示 [12], [14]。嵌入表主导了推荐模型的大小,并推动了模型规模的显著增长。尽管这种增长迅速,但它仍然受到内存容量需求相应增加的限制。为了进一步提高推荐模型的准确性,嵌入表的总容量已经增长到无法在单个服务器节点上支持的程度,这促使了分布式服务范式的需求。

III. 分布式推理

这是首个详细描述分布式神经网络推理的工作,针对深度推荐系统的独特目标和特性。由于这一初始实现更注重可部署性和可扩展性,而非性能与效率,因此在系统设计空间中存在许多优化机会,包括算法、软件、微架构、系统配置和任务映射等方面。因此,在本节中,我们为研究社区提供了通用系统设计和结构的概述。至少,神经推荐系统的分布式推理解决方案必须实现以下目标:(1)支持更多种类和更大数量的稀疏特征;(2)支持更大的嵌入表,即更大的嵌入维度、更少的哈希冲突和/或更少的压缩。

A. 分布式模型执行

假设深度学习模型可以表示为有向控制和数据流图。分布式模型简单地将神经网络划分为多个子网,每个子网可以独立运行。将每个子网独立对待提供了部署中的灵活性,因为用于服务模型的基础设施已经存在。图2b展示了这种划分的示例。传统上,这被称为模型或层并行。

1) 模型分片

我们将每个独立的子网称为一个分片(shard),而创建分片的过程称为分片化(sharding)。分片化在训练之前和之后都会进行。在训练之前,生成参数服务器分片以保存分布式训练的模型参数。在训练之后,模型发布期间,参数会根据预先的划分阶段从参数服务器重新分片并序列化到相应的推理分片。在此阶段,还会执行其他模型转换(例如量化),并且所有训练元数据都已可用于分片决策。训练后直接重新分片避免了重新加载和重新分片TB级模型所需的额外存储、计算和复杂性。

在图2b中,主分片执行所有稠密层,而任何划分的子网都被自定义的远程过程调用(RPC)操作符取代,这些操作符调用远程分片。只有稀疏操作符(如SLS)及其嵌入表被放置在远程分片上。这种启发式方法直接解决了嵌入表带来的内存限制,并保留了主分片上的计算密度。因此,远程分片也被称为稀疏分片。图3展示了这种方案的示例分布式跟踪,其中主分片(位于顶部)在每个网络中执行大部分计算。

图片名称

图3

由于嵌入表的巨大规模,单个表也可以被划分到多个分片上,如图2b所示。在这种情况下,稀疏特征ID会根据哈希函数拆分并发送到相应的RPC操作符。这是通过使用简单的模运算符将嵌入表行划分到多个分片上实现的。服务基础设施限制了分片之间不能存在图循环,因此每个分片都是无状态的,以避免进一步的复杂性。这一限制还在服务环境中提供了更大的灵活性,因为分片可能会失败并需要重启,或者可能需要添加副本。

2) 服务分片

分布式推理需要一个额外的特殊远程过程调用(RPC)操作符,它取代主分片中的子网并调用其相应的稀疏分片,如图2b所示。这种方案为大规模模型提供了直接的横向扩展支持。推理请求被发送到加载了主分片的服务器,当遇到RPC操作符时,会向相应的稀疏分片发出后续请求。所有分片上的推理服务由RPC服务处理程序(如Thrift或gRPC)和机器学习框架(如Caffe2、PyTorch或TensorFlow)组成 [30]–[32]。每个分片运行一个完整的服务处理程序和ML框架实例。

这种分布式架构支持与非分布式推理相同的复制基础设施。分片根据其负载和大规模部署的资源需求进行复制。例如,需要更多计算资源以满足QPS要求的分片将通过集群级管理程序在其集群上复制和部署副本。主分片和稀疏分片都运行完整服务堆栈的优势在于它们可以独立复制。每个单独的请求可以由与之前请求不同的机器组合处理,这进一步推动了无状态分片的需求。在我们的实验中未启用复制基础设施,因为使用了一组独立的服务器进行深入的每请求开销特征分析和分析。关于分布式推理对复制的预期影响和交互的讨论包含在第七节中。

B. 容量驱动的模型分片

最优的模型分片是一个具有挑战性的系统问题,因为配置数量众多且优化目标各异。虽然这一问题在训练背景下已被研究,但在推理背景下却带来了新的挑战。训练必须应对小批量大小和参数同步对准确性的影响,并旨在最大化单个实例上的吞吐量,该实例被输入大量数据 [33]。相比之下,数据中心推理必须在严格的延迟约束下应对不同的请求速率和大小,以及模型复制以满足这些需求的影响。此外,在推荐系统中,内存占用主要由计算稀疏的嵌入表主导。深度学习推荐推理的模型分片是为了支持大规模模型,因此我们将这一新挑战称为容量驱动的分片。

由于分片配置的数量众多,穷举搜索最优分片方案是不可行的。因此,我们使用启发式方法,这些方法依赖于模型架构,并包括对模型容量、计算资源和通信成本的测量或估计。简单来说,启发式方法应旨在最小化分片数量,因为更多的分片会消耗额外的计算和网络资源。我们在第六节中研究了这一假设的影响。启发式方法基于以下观察:
1) 稀疏层中使用的嵌入表占模型大小的绝大部分(> 97%)。
2) 稀疏层受内存和通信限制,而稠密层受计算限制。
3) 现有服务器基础设施无法支持稀疏层的内存需求,但可以支持稠密层的计算和延迟需求。
4) 深度推荐模型架构(如图2所示)可以并行执行稀疏操作符,这些操作符为后续的稠密层提供输入。

由于稠密层无法从分布式推理提供的额外计算或并行性中受益,因此我们的分片策略仅限于将稀疏层移动到远程分片。这是针对深度推荐模型架构的特定设计。图2显示,在嵌入表查找和池化之后,所有生成的嵌入向量和稠密特征在一系列特征交互层中组合。将交互层单独放置在自己的分片上会导致通信增加和计算密度降低。将该层与现有的稀疏分片放置在一起也会增加通信并违反无状态分片的约束。其他架构可能会从分片交互层中受益。例如,考虑一种模型架构,它为特定的嵌入表集提供了专用的特征交互层。这种架构可以将这些特征交互层与其相应的稀疏层分片,从而获得性能提升。尽管这是一种有趣的模型架构,但我们仅限于当前可用的模型,因此本工作选择专注于图2a中更传统的模型。

因此,分片稀疏层(1)直接解决了最大参数的容量问题,(2)有效地并行化了原本顺序执行的稀疏操作符,(3)通过隔离推理中受通信和计算限制的部分,实现了更好的资源分配。需要注意的是,推理传统上不是操作符并行的,因为操作符通常无法产生足够的工作量来抵消调度开销。额外的计算核心通过增加批处理级并行性来利用。

本工作评估了三种基于上述启发式的分片策略(表I)。这些策略应对了大规模深度学习推荐模型的新挑战:嵌入表的数量和大小各异、稀疏性的不均匀性以及实时服务环境的挑战。表I还展示了两种简单情况:(1)非分布式推理(单一分片)和(2)包含所有嵌入表的单一分片。

1) 容量均衡分片

一种直观的策略是将嵌入表均匀地分布在多个分片上。容量均衡分片确保每个稀疏分片具有相同的内存需求。这旨在最小化分片数量,目标是为单个服务模型实现最少的计算资源开销。

2) 负载均衡分片

每个稀疏特征表示为多热编码向量,该向量转换为嵌入表中的多索引查找和最终的池化操作。由于每个表的预期查找次数取决于特定特征及其在请求输入中出现的频率,容量均衡分片可能导致分片负载不均衡,从而执行显著不同的工作量。此外,查找次数与用于发送表索引的网络带宽成正比。这种总体不均衡可能导致某些分片成为关键路径瓶颈,与负载均衡分片配置相比,延迟会降低。

图3中展示的分布式跟踪示例说明了尽管网络延迟的不可预测变化也必须考虑,但一个稀疏分片可能会成为关键路径上的瓶颈。远程分片1和2是异步并行查询的。由于远程分片1执行了显著更多的工作,因此产生了更多的延迟开销。为了减少一个分片持续增加延迟的可能性,负载均衡策略根据嵌入表的池化因子(即预期查找次数)来放置嵌入表。池化因子通过从评估数据集中采样1000个请求并观察每个表的查找次数来估计。

3) 网络特定的装箱策略

推荐模型通常将用户特征和内容/产品特征分离到不同的网络中,以更有效地将用户-内容对批量处理 [1]。在本工作选择的模型中,用户网络的输出被输入到内容/产品网络中,因此它们必须顺序执行。图3展示了这一点,其中网络2依赖于网络1。当每个网络执行时,会调用单独的RPC操作符以访问其相应的嵌入表。如果表未按网络分组,则每个批次中可能会多次访问同一分片,以处理每个网络。这是不理想的,因为:1)在图3中,调用了三个RPC操作符,而如果按分片分组则只需调用两个(每个分片一个RPC);2)在数据中心环境中,服务器被复制以处理增加的请求时,无论哪个网络接收更多输入,两个网络的表都会被复制。

如果图3中远程分片1的服务器复制是由网络1的池化操作的高计算特性触发的,那么这种情况尤其不理想。网络2的所有嵌入表(可能被分片为数十GB)将消耗额外的、未充分利用的内存资源。为了解决这一问题并实现更高效的资源利用,我们评估了一种网络特定的装箱策略(NSBP),该策略首先按网络对表进行分组,然后根据给定的大小约束将它们装入“箱子”中。为了减少分片过程中消耗的数据中心资源,每个“箱子”最初是训练期间使用的现有稀疏参数服务器。如果参数服务器的“箱子”已满,则将其视为一个完整的分片。这减少了分片所需的网络带宽和计算编排。

C. 分布式推理实现

我们详细介绍了本工作中使用的定制开源框架。这为读者提供了一个具体的实现框架,以便理解我们的结果,并为使用其他框架实现提供指导。在本工作中,分布式推理构建在高度定制的Thrift和Caffe2变体之上,但该方法可推广到任何RPC或ML框架 [30], [31]。尽管PyTorch已经吸收并取代了Caffe2作为未来的最先进框架,但本工作中使用的Caffe2基础设施在两者之间是共享的。Thrift通过加载模型并将接收到的请求适当地拆分为推理批次来处理排名请求。我们使用了一个修改版的Caffe2,其中包含了发出Thrift请求的RPC操作符。对稀疏分片的中间请求通过通用服务发现协议进行路由。所有服务器间的通信通过以太网上的标准TCP/IP协议栈进行。模型在训练后被转换为分布式推理格式。一个自定义的分区工具使用用户提供的配置来分组嵌入表及其操作符,插入RPC操作符,生成新的Caffe2网络,然后将模型序列化到存储中。

IV. 跨层ML操作符和通信感知的特征分析

分布式推理的性能由系统多个层次的选择决定,特别是数据中心服务层(调度、发现和网络)、机器学习框架层以及机器学习操作符本身。跨这些层次测量工作负载对于理解开销、归因成本、优化目标组件以及做出高层系统设计决策非常重要。由于没有现成的分析工具可以执行这种跨层特征分析,我们构建了一个自定义的跨层分布式跟踪框架来测量分布式推理工作负载。

A. 通过分布式跟踪捕获工作负载

分布式跟踪是一种经过验证的技术,用于调试和调查跨依赖服务的性能问题 [34], [35]。我们利用这种技术来研究分布式推理基础设施的性能影响。在生产服务堆栈的多个层次中添加了自定义检测,以提供请求/响应序列化、RPC服务样板设置以及每个请求的模型执行的完整成本视图。

Thrift 和 Caffe2 都提供了检测钩子(instrumentation hooks)来捕获执行中的关键点,而服务处理程序则需要更侵入式的修改,尽管仍然保持轻量级。我们选择了源代码检测而非其他二进制检测方法,因为它的开销更小,并且可以与 Thrift 的

1
RequestContext
抽象接口结合,以传播分布式跟踪的上下文数据。在每个跟踪点,特定于层的元数据和挂钟时间戳(wall-clock timestamp)会被记录到一个无锁缓冲区中,然后异步刷新到磁盘。挂钟时间是首选,因为它的顺序有助于实现有用的跟踪可视化。此外,大多数时间跨度(span)较小且是顺序的,这使得挂钟时间可以作为 CPU 时间的代理。每个请求的每个分片的 CPU 时间也会被记录下来,以验证这一说法。跟踪点随后会被收集并在离线状态下进行后处理,用于开销分析并重建事件的可视化,类似于图3。

图3展示了一个分布式跟踪的示例。分片通过水平切片分隔,主分片位于顶部处理请求。该示例跟踪表示单个批次的执行,但根据排名请求的数量和可配置的批次大小,通常会并行执行多个批次。此外,操作符的执行顺序是本工作中研究的推荐模型的典型顺序。操作符被调度为顺序执行——除非像 RPC 操作符那样明确是异步的——因为其他核心通过请求级和批次级并行性来利用。首先执行初始的稠密层,然后是稀疏查找,接着是后续的特征交互层和顶部的稠密层。

如第III-A节所述,每个稀疏分片都是一个完整的RPC服务,以实现部署的灵活性,因此与本地内联计算相比,会带来与服务相关的开销。图3展示了每个稀疏分片中未用于执行SLS操作符的时间即为延迟开销。在非分布式情况下,SLS操作符直接在主分片中执行。额外的延迟归因于网络链路、输入和输出的额外(反)序列化,以及准备和调度Caffe2网络所花费的时间。尽管存在这些开销,我们也看到分布式计算的异步特性使得稀疏操作符能够实现更多的并行化。这一特征分析的一个重要启示是,在增加并行化以减少延迟开销和减少并行化以降低计算和数据中心资源开销之间存在权衡。此外,这种权衡是特定于每个模型的。

B. 跨层归因

与简单的基于操作符的计算或端到端(E2E)延迟计数器相比,捕获跨层跟踪提供了计算和延迟的整体视图。图3中的示例流程展示了RPC序列化和反序列化以及软件基础设施的请求/响应处理带来的计算开销。每个线程在调度和记录异步RPC操作符时也会产生额外的计算。简单的端到端延迟开销可以轻松地在主分片中测量。然而,由于每个稀疏分片是并行执行的,延迟开销的归因更为复杂,涉及稀疏分片请求之间的重叠。为了简化分析,每个主分片请求中最慢的异步稀疏分片请求被用于延迟分解。使用相关稀疏分片中的网络、序列化和RPC服务延迟。由于不同服务器上的时钟可能存在偏差,网络延迟通过主分片中测量的未完成请求与稀疏分片中测量的端到端服务延迟之间的差异来计算。

V. 实验方法与工作负载

我们评估了生产级的分布式服务软件基础设施,以量化其对计算需求、延迟和分片效果的影响。本研究的目的是评估分布式推理作为支持大规模深度推荐模型的一种手段的实用性,并为这一新颖工作负载的进一步系统探索提供基础。在本节中,我们描述了为本研究选择的数据中心深度学习推荐模型以及底层硬件平台。

A. 推荐模型

本研究使用了三个类似DLRM的模型,分别称为DRM1、DRM2和DRM3,这些模型涵盖了一系列大型模型属性,例如不同的输入特征和嵌入表特性。前缀“D”用于区分这些分布式模型与最近深度学习推荐工作中讨论的特定模型 [1]–[3]。DRM*模型是众多可能的、不断演变的模型的一个子集,并因其稀疏特征特性而被选为分布式推理的早期候选。研究这些模型的目的是通过我们的跨层分布式跟踪分析为评估分布式推理的开销提供基础,它们不应被解释为像MLPerf基准测试套件 [9] 中包含的典型基准模型。

图4、图5和表II展示了大型神经推荐模型中的多样性。DRM1和DRM2有两个网络,每个网络都有其各自的稀疏特征,而DRM3只有一个网络。由于使用了每个模型的真实采样请求,推理请求的大小以及相应的计算和延迟在模型之间也有所不同。所有参数均以单精度浮点数形式存储,未进行压缩。第VII-D节讨论了压缩的影响。图4展示了每个模型的每操作符组归因,这是非分布式模型所有采样请求的简单平均值。DRM1和DRM2是最相似的架构,这反映在图4中。与DRM3相比,DRM1和DRM2具有更复杂的结构,这体现在额外的张量变换成本上。与本工作更相关的是,稀疏操作符在所有操作符中所占的比例远高于DRM3。具体而言,稀疏操作符在DRM1、DRM2和DRM3中分别占所有操作符时间的9.7%、9.6%和3.1%。尽管它们的计算比例较低,但稀疏操作符在DRM1和DRM2中占模型容量的>97%,在DRM3中占>99.9%。

图片名称

图4

为了将所有模型适配到单个256GB服务器上,大于给定阈值的嵌入表按比例缩小。这为表I中列出的所有分片策略(包括单一分片)提供了计算和延迟开销的直接比较。原始数据中心规模的模型要大得多。

图5展示了每个模型中嵌入表大小的分布。DRM1的大小为200GB,包含257个嵌入表,最大的表为3.6GB。DRM2的大小为138GB,包含133个嵌入表,最大的表为6.7GB。最后,DRM3的大小为200GB,包含39个嵌入表,最大的表为178.8GB。与DRM3相比,DRM1和DRM2展示了嵌入表大小的长尾分布,这解释了图4中显示的额外稀疏操作符成本。相比之下,DRM3由一个单一的大表主导。这些嵌入表是先前深度学习推荐工作中讨论的表的代表性版本 [1], [3], [9]。

图片名称

图5

对于DRM1和DRM2,评估了十种分片配置。表II描述了DRM1的分片策略结果。在负载均衡配置中,每个分片的容量差异高达50%,而在容量均衡配置中,每个分片的总容量相同。在容量均衡配置中,每个分片的估计负载差异高达371%(在8个分片的情况下,分片4和分片3之间),而在负载均衡配置中,每个分片的总池化因子相同。最后,网络特定的装箱策略(NSBP)将每个分片的表限制为单一网络。这在2分片配置中最为明显,其中每个网络被放置在自己的分片上。分片2消耗的内存容量是分片1的4.75倍,但估计仅执行分片1计算工作的6.3%。

由于分片超大表的技术挑战,DRM3仅使用NSBP策略进行分片,而未采用容量均衡或负载均衡策略。由于它由一个单一的大表主导,每增加一个分片,最大的表会进一步拆分,而较小的表则分组为一个分片。主导表的池化因子为1,因此只有跨越该表的一个分片会被访问。例如,给定四个稀疏分片,最大的表被划分为三个分片,其余的表分组为一个分片。每次推理仅会访问两个分片:一个用于分片的大表,另一个用于较小的表。

B. 测试平台

为了进行特征分析,我们使用了代表数据中心环境的两类服务器。SC-Large代表数据中心中的典型大型服务器,配备256GB DRAM和两个20核Intel CPU。SC-Small代表典型的更高效的Web服务器,配备64GB DRAM和两个时钟频率较低的18核Intel CPU,并且网络带宽低于SC-Large。由于SC-Small的DRAM容量有限,只能在该平台上测试部分配置。第六节中讨论的大部分结果是在SC-Large平台上收集的,以便与非分布式模型进行公平比较。关于服务器平台影响的讨论见第七节。

所有推理实验均在CPU平台上运行,如第二节所述。预留的裸金属服务器与生产推荐排名服务器位于同一数据中心,并使用相同的内部网络。它们在数据中心内的位置代表了典型的推理服务层。分片被分配到独立的服务器上,而不是共置。分片到服务器的映射在重复试验中随机分配。推荐排名请求从生产服务器中采样,并在测试基础设施上重放;去标识化的请求数据库在五天时间内均匀采样,以捕获请求中的昼夜行为。生产重放器在将请求发送到推理服务器之前对其进行预处理和缓存,使用与在线生产排名相同的网络基础设施。

在大多数实验运行中,推理的批次大小设置为生产默认值,其中每个批次代表要排名的推荐项目数量,并并行执行。在第六节中,请求被串行发送,以隔离固有开销。在第七节A部分中,请求被异步发送,以模拟更接近生产环境的更高QPS率。

VI. 分布式推理分片分析

在本节中,我们讨论了跨层分布式跟踪分析的结果。我们发现,在严格的延迟目标和数据中心环境的计算规模下,设计空间比简单地最小化分片数量更为复杂。我们的主要结论如下:

  • 增加分片数量可以通过有效增加模型并行性来管理RPC操作符带来的延迟开销。
  • 然而,增加分片也会因服务样板和调度而带来额外的计算成本。
  • 串行发送的阻塞请求在分布式推理中的P50、P90和P99延迟表现总是更差,这是由于简单的Amdahl定律限制。深度推荐系统的嵌入查找在端到端延迟中占比不够大,无法从增加的并行性中受益。
  • 与容量均衡策略相比,负载均衡策略对端到端延迟的影响不显著。网络特定的装箱策略(NSBP)在延迟方面表现最差,但在计算方面表现最佳。
  • 某些模型的工作量不足以并行化,因此无法从额外的分片中减少开销。

研究结果表明,分布式推理是服务大规模深度学习推荐模型的实用解决方案,而现有的服务范式无法支持这些模型。研究结果还暴露了对分片之间网络链路和支持通信管理的软件基础设施进行针对性分析的需求,因为稀疏分片上执行的工作量较少。结果表明,自动分片方法是可行的,但需要足够的分析数据,因为嵌入表行为的多样性和分片策略的复杂权衡。

A. 分布式推理开销

图6和图7展示了所有模型在串行阻塞请求下,所有分片策略的P50、P90和P99端到端(E2E)延迟和总计算时间开销。每种分片策略的描述见表I。请注意,单一分片是非分布式情况,而1分片是所有嵌入表位于一个稀疏分片上的最坏情况。P90和P99延迟是推理服务的典型指标,因为如果推理请求未在SLA目标内处理完成,则会返回准确性较低的备用推荐结果。为了完整性,P50也被展示出来,以显示中位情况,而不是平均值,因为先前的工作中讨论了长尾延迟 [1]。

图片名称

图6

图片名称

图7

B. 延迟开销

1) 增加分片可以减少延迟开销:所有模型的分布式推理配置均出现了端到端延迟增加,如图6和图7所示。稀疏操作符的并行化与异步RPC操作符的结合不足以抵消增加的网络延迟和软件层开销。对于DRM1和DRM2,延迟表现最差的配置是一个稀疏分片,这是不切实际的最坏情况,所有嵌入表都放置在一个分片上,没有并行化工作。令人鼓舞的是,在2分片负载均衡配置下,DRM1的P99延迟开销仅为7.3%。在8分片负载均衡配置下,这一开销降至P99的1%,P50中位情况的11%。这表明,通过简单的分片策略可以实现最小且实用的延迟开销。出乎意料的是,2分片NSBP策略在DRM1的P99延迟中表现最差,在DRM2中几乎表现最差。回顾表II,大部分工作(池化因子)被分配到一个分片上,因此在P99情况下,2分片NSBP配置实际上类似于1分片配置的边界情况。

2) 恒定开销最终占据主导地位:随着分片数量的增加,每个分片的工作量减少,网络延迟和额外的软件层开销占据主导地位,如图8所示。网络延迟通过主分片的未完成请求时间与稀疏分片的总端到端时间之间的差异来测量。这个时间包括内核数据包处理和转发时间。对于所有分布式推理配置,网络延迟均大于操作符延迟。分布式推理总是会损害这些模型的延迟。换句话说,如果稀疏操作符平均产生足够的工作量,那么模型将适合分布式推理。并且,在稀疏操作符工作量足够的情况下,延迟可以得到改善。这为系统架构师、模型架构师和特征工程师提供了一个多学科合作的机会,以平衡模型资源消耗、性能和准确性。

3) 分片影响取决于模型架构:对于DRM3,分片数量没有显著影响。DRM3由一个单一的大表主导,如图5所示,该表在分片之间拆分。即使在8分片的情况下,每个推理请求也只访问2个分片——一个分片包含较小的表,另一个分片包含分片后的最大表的条目,如图11a所示。

4) 深入的延迟层归因:图8展示了跨跟踪层的P50(中位数)延迟的细分。图8a显示,只有工作负载的嵌入部分在不同分片策略下显著受到影响,这是意料之中的,因为这部分工作负载被卸载到稀疏分片上。Dense Ops是指所有非嵌入表查找和池化操作的ML操作符;Embedded Portion是指所有嵌入表查找和池化操作符。对于单一分片配置,这是操作符本身的时间,而对于分布式推理配置,这是等待稀疏分片响应的时间,如图8a所示。RPC Serde指所有序列化和反序列化请求时间,而RPC Service是指严格未花费在Caffe2网络或序列化/反序列化上的其他时间。最后,Net Overhead是指网络中未用于执行操作符的时间,例如异步操作符的调度。对于分布式推理配置,嵌入部分延迟的影响反映了图6和图7中显示的开销。

图片名称

图8

在DRM1的单一分片配置中,嵌入部分仅占延迟的约10%,而在单分片配置中,这一比例为32%。在8分片负载均衡的最佳分布式推理情况下,嵌入部分占总延迟的15.6%。相比之下,DRM3的嵌入部分在分片增加时没有显著变化,因为只有主导的大表被进一步分区。DRM3的延迟和计算变化归因于缓存效应和与更多服务器节点通信的网络变异性。图8b进一步归因了嵌入部分内的延迟——每个条形堆栈代表图8a中的嵌入部分。对于P99,嵌入部分的重要性降低,而主分片上的稠密操作符和RPC反序列化开始占据主导地位,这是由于非常大的推理请求大小。这就是为什么P99延迟开销比P50更有利的原因。

C. 计算开销

理解计算开销对于最小化数据中心规模的额外资源需求至关重要。分片策略是系统设计者平衡延迟约束(如前一节所述)和影响资源需求的计算开销的一种方法。

1) 增加计算开销是减少延迟开销的权衡:高计算开销是灵活、易于部署系统的代价。如前一节所述,增加稀疏分片可以减少分布式推理的延迟开销。然而,计算开销也会增加,因为每个分片都会调用完整的Thrift服务。图9显示,对于所有模型,分布式推理总是会增加计算开销,因为需要额外的RPC操作。更重要的是,图9表明计算开销与RPC操作的数量成正比。NSBP策略的计算开销最小,因为它执行的RPC操作最少。回想一下,NSBP策略限制每个分片不混合来自不同网络的嵌入表,因此每个分片在每次推理中仅被调用一次。对于多个网络,每个网络的并行化程度较低。相比之下,其他分片策略可能会更多地并行化每个网络,从而调用更多的RPC操作,导致整体计算开销增加。

图片名称

图9

2) 计算开销影响数据中心资源:当计算开销发生在主分片上时,问题尤其严重,因为它会增加计算驱动的复制和资源需求,以处理相同的QPS。当主分片上发出RPC操作所需的计算开销超过通过卸载嵌入部分节省的计算时,就会发生这种情况。对于具有许多大嵌入表和低池化因子的模型架构,这种情况更有可能发生。研究结果为探索这些拐点提供了动力,这些拐点应作为未来自动分片方法的输入,并依赖于模型属性和软件基础设施。这一点在第七节中进一步讨论。

D. 分片策略效果

前一节中已经确定了由于分片数量增加而导致的延迟增加或计算开销增加之间的权衡。分片数量是系统设计者平衡计算开销与延迟约束的一个直接调节手段。本节讨论的分片策略为设计者提供了另一个调节手段,但其对延迟和计算的影响更为微妙。

1) NSBP是最具可扩展性的策略:负载均衡和容量均衡配置之间的延迟开销没有显著差异。然而,网络特定的装箱策略(NSBP)因受额外分片影响较小而表现出不同。这一观察的最重要结论是,NSBP是评估中最具可扩展性的策略,因为它调用的RPC操作较少。

NSBP展示了最不平衡的每分片延迟。回想一下,在这种策略中,嵌入表首先按网络分组,然后根据网络和大小分配到分片。来自不同网络的表永远不会分配到同一分片。图10通过负载均衡和NSBP策略的每网络操作符延迟更清楚地展示了这一点。分片1和分片2包含执行最多工作但表大小最小的第一个网络,如图10b和表II所示。对于NSBP,这对延迟产生了负面影响,因为并行化的工作较少,但计算开销由于较少的并行化而受到较小的影响,因此调度和服务开销也较少。请注意,Net1和Net2是顺序执行的,因此它们对端到端延迟的累积影响是相加的。NSBP策略在资源利用率方面的优势将在第七节进一步讨论。

图片名称

图10

2) 负载均衡和容量均衡策略之间的开销差异很小:图12展示了DRM1在所有请求下的8分片配置的每分片操作符延迟;DRM2显示出类似的趋势,为简洁起见省略。回想一下第III-B节,负载均衡策略预计通过消除任何分片成为关键路径来降低延迟开销。然而,与端到端延迟相比,两种策略的每分片操作符延迟都微不足道。此外,负载均衡和容量均衡策略之间的延迟差异并不显著,正如表II中估计的池化因子所暗示的那样。在这种规模下,池化因子太小,无法显示出明显的效果。对于DRM1和DRM2,负载均衡和容量均衡分片策略之间的最大影响来自分片数量的增加。

E. 模型多样性

模型属性也会影响分布式推理性能。网络数量、嵌入表数量、大小分布和相应的池化因子是最相关的模型属性,并在第V-A节中对评估模型进行了描述。与DRM1和DRM2中的长尾嵌入表相比,DRM3的稀疏操作符总计算量较少,并且由一个单一的大嵌入表主导。DRM1和DRM2的延迟受益于增加分片,但DRM3没有显示出这种优势,如图7和图8所示。

1) 分片效果取决于模型架构:图11单独展示了DRM3的每分片操作符延迟和嵌入部分细分,以显示增加分片不会改善延迟。主要原因有两个。(1) 增加分片只会分区一个容量主导的嵌入表,这不会并行化任何显著的计算。图11a显示分片1执行了大部分计算,因为它包含除最大表之外的所有嵌入表,而最大表在其他分片之间分区。(2) 即使较小的表被分区,相对较低的计算量仍然无法实现实际的延迟改进,因为网络延迟在端到端延迟开销中占据主导地位。因此,我们得出结论,只有像DRM1和DRM2这样具有长尾嵌入表和较高池化因子的模型才能从分片中受益。

图片名称

图11

F. 批处理效果

推理的批处理大小将请求拆分为并行任务,是在满足SLA和QPS目标的前提下平衡吞吐量和延迟的关键。为了展示其与分布式推理的交互,我们将批处理大小人为设置为每个请求一个批次。较小的批处理大小增加了每个请求的任务级并行性,可以减少延迟,但会增加任务级开销并减少数据并行性,从而可能降低吞吐量。相比之下,较大的批次增加了稀疏操作符的工作量,并受益于分布式推理的并行化。深度学习推荐推理的批处理大小是一个正在进行的研究课题 [2]。

1) 分布式推理在足够大的批处理大小时可以改善延迟:图13显示,在使用8分片容量均衡或负载均衡配置时,分布式推理可以在DRM1单批次情况下改善延迟。这是因为稀疏操作符执行了足够的工作,从而充分受益于并行化,如图13b所示。DRM2显示出类似的趋势,但由于请求较小,影响不那么显著。在这种情况下,较大的批次可以被视为具有较大池化因子的嵌入表的代理,其显著特征是通过RPC发送的额外查找索引增加了稀疏操作符的工作量和网络需求。DRM3未显示,因为其请求通常足够小,默认批处理大小下每个请求只有一个批次。

图片名称

图13

2) 批处理大小可以管理分布式推理的计算开销:图14强调了计算开销的倍增效应——每个额外的批次都会发出相应的RPC操作,从而增加计算需求。例如,由于DRM1的NSBP策略每个分片发出一个RPC,其计算开销随着分片的增加而增加的速度比负载均衡策略慢。在每请求一个批次的情况下,分片带来的计算开销的边际增加不那么严重。因此,在探索分布式推理计算开销时,必须考虑批处理大小。

图片名称

图14

VII. 数据中心环境中的影响

在本节中,我们讨论了分布式推理对深度推荐系统在数据中心中的影响。分析延迟和计算开销对于理解分布式推理对SLA目标和额外计算资源的影响非常重要。然而,第六节中的分析集中于一个简化的场景,以归因每个请求的开销,其中(1)每个请求被串行处理,(2)服务器具有相同的SC-Large硬件配置,这些配置是过度供应的。为了模拟更具代表性的服务环境,我们在DRM1上进行了两个额外的实验,DRM1是计算最密集的模型。首先,我们以更高的25 QPS速率发送请求,覆盖所有分片策略,并使用相同的SC-Large硬件配置。其次,我们在更典型的Web服务平台SC-Small上重新运行负载均衡配置,与SC-Large进行比较,请求再次被串行发送。我们将结果置于一个服务环境中,其中模型实例被复制以处理实时流量。最后,我们讨论了当前在大型数据中心规模的深度推荐系统中实施的现有压缩技术。我们的主要观察结果如下:

  • 以更高的QPS发送的请求(代表数据中心环境)在P99延迟下表现更好,因为资源可用性得到了改善。
  • 在SC-Large和SC-Small之间,稀疏分片的每请求延迟没有显著差异,这为通过低功耗的稀疏分片服务提供了提高效率的机会。
  • 分片复制通过独立分配模型的稠密和稀疏部分的资源,提供了提高服务效率的机会。
  • 压缩是分布式推理的补充,本身无法解决模型的可扩展性问题。

A. 高QPS环境

请求重放器被配置为以25 QPS的速率向DRM1的主服务器发送请求,DRM1是所有评估模型中计算最密集的。每个DRM1配置的开销图与单一配置相比,如图16所示。在25 QPS实验中,所有配置的开销都比串行发送时的相同配置低,如图6所示。在几乎所有配置中,P99延迟都比单一配置有所改善,如图6所示。

此外,在某些配置中,例如8分片的容量均衡和负载均衡配置,P50延迟比单一配置有所改善,或者开销小于约0.05%。在具有足够稀疏操作符计算量的模型中,分布式推理在高QPS场景中具有更好的延迟特性。

B. 稀疏分片平台的效率

为了对分布式推理中增加服务器进行公平比较,稀疏分片使用了相同的硬件平台。然而,图8和图12显示,稀疏分片的计算需求远低于主分片,因为本工作中的分片是容量驱动的,而非计算驱动的。因此,我们使用较轻量级的SC-Small服务器运行了DRM1的额外配置,DRM1仍然是计算最密集的模型。图15显示,当在原始较重量级的SC-Large和较轻量级的SC-Small服务器上运行时,每分片操作符延迟几乎相同。回想一下,SC-Large服务器具有更多且更快的核心,以及4倍的内存容量,导致能耗增加。这表明,通过粗粒度的平台专业化,稀疏分片可以在集群级别提高服务效率和能源效率。

图片名称

图12

图片名称

图15

C. 数据中心中的复制

在本节中,我们简要讨论了分布式推理如何改善分片复制。支持具有数百GB内存占用的单一模型需要配置效率低下的服务器。例如,大部分内存容量将用于大型嵌入表,但大多数核心将用于使用显著较小的稠密参数。对于DRM1、DRM2和DRM3,大多数计算仅涉及模型内存占用的不到3%。这种低效性在动态复制推理服务器以支持数百万或数十亿用户的QPS时会被放大。在这种情况下,稠密层带来的大负载(如图4所示)将导致整个模型被复制到额外的服务器上,包括所有嵌入表。分布式推理通过允许为稠密层分配计算资源,为稀疏层分配内存资源,缓解了这种低效性。换句话说,复制的内存需求减少了。我们将模型的稠密层和稀疏层分片为独立组件的启发式方法简化了这种分配。最后,分片策略在复制中起到辅助作用。虽然分片稀疏操作符的决定对减少复制的内存需求影响最大,但表II显示分片策略可以进一步影响复制。回想一下,DRM1由两个主要网络组成,其中Net 2消耗的内存资源是Net 1的4.75倍,但计算资源仅为Net 1的6.3%。NSBP分片策略将每个分片限制为仅包含Net 1或Net 2,因此与网络无关的分片策略相比,其并行化效果较差。然而,NSBP可以通过将计算最密集的嵌入表分组在一起,进一步提高资源利用率,即使是在需要更多分片的大型模型中。这种在改善延迟和提高资源利用率之间的分片策略权衡需要基于每个模型的具体情况进行,并鼓励进一步研究分片自动化。

D. 模型压缩

模型压缩是限制模型容量的传统方法。我们在表III中展示了DRM1的压缩模型大小作为参考。当前数据中心模型中部署的所有压缩技术(包括量化和剪枝 [22])被用于生成压缩模型。所有表均按行线性量化至至少8位,足够大的表被量化至4位。表根据模型架构师指定的阈值幅度或训练更新频率手动剪枝。量化和剪枝选项的选择是为了保持准确性和延迟。我们注意到,压缩后延迟和计算略有改善,但由于这不是本工作的重点,我们将此分析留待未来工作。我们推测原因是内存局部性的相对改善。更相关的是,表III显示压缩模型缩小了5.56倍。虽然显著,但即使有这些节省,大型模型仍然无法适配到一台、两台甚至四台配置了约50GB可用DRAM的商品服务器上。因此,仅靠压缩不足以支持新兴的大型深度推荐系统。

参考