介绍

kuaishou在2020《Kraken: Memory-Efficient Continual Learning for Large-Scale Real-Time Recommendations》提出了它们的实时推荐系统。

摘要

现代工业推荐系统常常使用深度学习(DL)模型,这些模型通过更多的数据和模型参数来实现更高的模型准确性。然而,当前的开源DL框架,如TensorFlow和PyTorch,在训练具有数太字节参数的推荐模型方面显示出相对较低的可扩展性。为了有效地从每天生成数百太字节(terabytes)训练数据的数据流中学习大规模推荐模型,我们引入了一个名为Kraken的持续学习系统。Kraken包含一个特殊的参数服务器实现,它能够动态适应持续变化的sparse特征集,以持续训练和服务推荐模型。Kraken提供了一个感知sparse性的训练系统,该系统对dense和sparse参数使用不同的学习优化器,以减少内存开销。使用真实世界数据集进行的广泛实验证实了Kraken的有效性和可扩展性。Kraken可以在相同的内存资源下提高推荐任务的准确性,或者在保持模型性能的同时将内存使用量减少到原来的三分之一。

一、引言

近年来,推荐系统已经成为许多流行移动应用的基石。它们为各种内容(包括新闻文章、短视频和广告)生成个性化排序,从而提高用户与这些应用的互动体验。正如流行商业分析师所报告的,推荐系统通过增加用户参与度,为许多大公司如亚马逊和Facebook带来了相当一部分收入[1]-[3]。

时间敏感性对于推荐系统实现合理性能至关重要。例如,用户在与移动应用互动时的兴趣通常是非常非平稳的、季节性的,并且对趋势敏感。这被称为概念漂移[4]。推荐系统的另一个重要问题是所谓的冷启动问题[5],即在有限的时间内推断新用户的偏好或新item潜在的受众。解决这些问题的一个常见方法是:使用实时持续学习(或在线学习)[6]、[7],这意味着不断用新数据训练推荐模型以保持模型的新鲜度。这种策略适用于许多经典机器学习模型,如逻辑回归[8]、[9]和矩阵分解[10]。然而,随着深度学习(DL)在推荐系统中的兴起,DL模型的在线学习在系统可扩展性和模型质量方面面临挑战。

与用于计算机视觉(CV)和自然语言处理(NLP)的经典机器学习模型或DL模型不同,推荐系统的DL模型使用大量sparse分类特征,这些特征表示为一维或多维的高维独热二进制向量。随着不同sparse特征的数量达到数百万甚至数十亿以提高准确性,模型大小变得多太字节,因此不适合单个GPU甚至单个服务器的内存。此外,与有限的内存资源相比,每分钟都有大量新生成的内容和用户行为需要被表示到DL模型中。庞大的模型和不断流动的数据流为训练系统创造了极高的内存压力。在通过大量数据训练巨型模型的压力下,更难以有效地服务实时更新的模型。

现有系统不足以克服这些挑战。一般的开源DL框架,如TensorFlow[11]和PyTorch[12],对于CV和NLP领域批量训练复杂DL模型进行了高度优化。先前在生产环境中的研究[1]、[13]、[14]表明,这些通用DL框架由于不善于处理大规模sparse特征,因此在大型推荐模型方面扩展性不佳。此外,不足的端到端在线学习支持使它们无法在不断增长和持续更新的模型下工作。即使是一些开源框架的内部版本也被认为需要高维护成本才能启用在线学习[15]。

在本文中,我们介绍了Kraken,这是一个考虑到sparse嵌入的生产就绪系统,用于优化在线学习和服务大规模推荐模型。据我们所知,这是第一份包含构建工业级推荐系统大规模持续学习系统在系统和算法方面的足够细节的论文。Kraken的核心是一个感知sparse性的训练系统,它有一个特定的参数服务器实现,结合数据并行和模型并行来训练推荐模型。专门的参数服务器支持自动特征准入和过期机制,以高效管理在线学习期间sparse嵌入的生命周期,利用有限的内存资源以获得更好的模型性能。此外,Kraken的在线服务系统将sparse嵌入的存储和模型预测的计算解耦,这显著节省了网络和存储成本。我们在TensorFlow 1.14之上实现了Kraken的训练系统,因此Kraken与TensorFlow API兼容。我们通过使用真实世界数据集进行离线实验和在线A/B测试来检验Kraken。结果揭示了与原始TensorFlow系统相比,Kraken的有效性和可扩展性。

二、背景与动机

A. 推荐系统中的深度学习

图1a展示了推荐系统的概览。对于给定的用户查询u,带有各种用户、item和上下文信息,推荐系统通常执行两步程序来从可能包含数十亿item的数据库中生成最相关items的$\lbrace x_i \rbrace$的排y序列表。

第一步称为检索(retrieval),它通过使用简单的机器学习模型和轻量级人为定义的规则的组合来返回数百个item,以提高效率。

在缩小候选池(candidate pool)后,后续的排序(ranking)步骤根据复杂深度学习模型生成的排序分数对所有item进行排序。排序分数通常是$ P(y \mid u, x_i)$的估计,即用户u在查看item $x_i$后,后续用户行为标签y(例如,点击、点赞)的概率。

为了训练这些模型,真实的用户反馈数据以及用户和上下文信息被记录在日志中作为训练数据。为了获得排序分数的准确预测,现代推荐模型使用大量sparse特征和复杂的DL模型[16]-[18]。sparse类别型特征通常用于表示任何两个一般实体之间的交互。例如,用户偏好特征可以是用户曾经点击过的最后K个视频ID的列表。为了有效处理sparse特征,推荐模型经常使用一种称为sparse嵌入的技术将它们转换为低维dense表示。如图1b所示,sparse特征可以被视为sparseID的向量。每个sparse特征与一个嵌入表配对,并且在转换过程中,特征的每个sparseID用于查找嵌入表中的唯一列(称为嵌入向量)。然后,通过逐元素收集操作(称为池化操作)将所需的嵌入向量组合成一个dense向量。新形成的dense向量成为模型其余部分的输入,以进行最终预测。

图片名称

图1 推荐系统概述。(a) 推荐系统中的两个步骤,检索和排名。(b) 推荐系统的典型深度神经网络(DNN)模型架构。sparse特征(一系列ID,例如UID(用户ID)和VID(视频ID))首先会被映射到不同嵌入表中的dense嵌入向量(模型的sparse部分)。然后,所需的嵌入向量被组合并成为模型其余部分的输入,以进行最终预测。(c) 显示了嵌入表中的哈希冲突。

这些sparse嵌入表通常被称为推荐模型的sparse部分(sparse part)。模型的其余部分,包括多层全连接深度神经网络(DNN),被称为dense部分(dense part)。在数据大小和访问模式方面,sparse和dense部分之间存在巨大差异。如表I所示,sparse部分的大小可能是dense部分的1000倍甚至更大。然而,在每次训练或预测的小批量中,只有有限数量的嵌入向量在sparse部分被访问,而dense部分每个批次都被完全访问。Kraken重新设计了sparse嵌入的存储,并采用了更好的哈希策略,允许它支持更大的嵌入表。

当前的开源DL框架,如TensorFlow和PyTorch,使用dense可变矩阵(固定大小数组)来表示sparse嵌入。对于这些框架的常见做法是,这些sparseID需要被哈希到一个预定义的、可计数的集合中,以控制每个嵌入表的大小,称为哈希技巧(见图1c)。例如,对于一个ID为j的视频,其嵌入存储在大小为M的数组的索引(hash(j) mod M)处。这种方法可能很棘手,因为一些sparseID比其他ID更频繁地被访问(例如,那些视频更受欢迎,被更多用户点击)。不幸的是,当热门ID被哈希到同一个哈希桶时,由于值重叠,会导致预测准确性降低。避免哈希冲突的一个简单方法是增加哈希表大小,这会浪费未使用哈希桶的内存。Kraken重新设计了sparse嵌入的存储,并采用了更好的策略,允许它支持弹性扩展的嵌入表

B. 推荐模型的并行范式

随着机器学习模型获得更多参数,由于其有限的计算和内存资源,单台机器不足以训练和服务大规模模型。许多先前的研究[11]、[19]-[21]提出了用于CV和NLP中使用的大规模神经网络的并行训练系统。典型的并行机制包括数据并行和模型并行[20]。数据并行将训练数据划分为多个数据集,每个worker一个,而模型并行将模型划分为可以并行训练的多个部分。

作为常见做法,推荐模型的并行训练结合了模型并行和数据并行,这是由于不同类型的参数特征(图2)。

  • 对于模型的sparse部分:嵌入表是模型并行的,并通过哈希在多个worker之间共享
  • 而对于模型的dense部分:DNN是数据并行的,每个worker有一个独特的副本

图片名称

图2 推荐模型的典型并行范式结合了模型并行性和数据并行性

sparse参数的大内存消耗基本上决定了我们至少需要多少worker来完成训练和服务。Kraken提出了几项优化,以在最低的计算资源开销下提供内存高效的学习和服务。

C. 大规模在线学习的必要性和挑战

先前的研究工作表明,在许多任务中增加深度学习模型的大小可以极大地提高它们的准确性和预测能力,而不会过拟合[22]、[23]。这一观察也适用于推荐系统,因为更大的嵌入表可以捕捉更细粒度的用户行为和item属性。图3a展示了基于DNN的推荐模型在三个工业数据集上的性能随着模型大小的增长而显著提高[13]、[24]。在大规模学习推荐模型方面,在线学习比批量训练有许多优势。

  • 首先,工业推荐系统通常每天收集高达数百太字节(terabytes)的训练数据。在线学习通过流式实例使高效训练成为可能,每个训练实例只需要处理一次。
  • 其次,在线学习允许新模型更频繁地更新并部署,比批量训练更重要,这对于解决第I节中提到的问题,如冷启动和概念漂移问题。

为了展示在线学习带来的好处,我们运行了一个在线A/B测试来比较Kraken的两种模式:

  • 一种使用每五分钟更新一次模型的在线学习
  • 另一种使用在测试期间不更新的静态模型

图3b绘制了A/B测试期间具有不同设置的两种模型的平均AUC,其中在线学习模型保持模型准确性稳定,但静态模型的AUC在一小时后下降了4.7%。它得出结论,在线学习在跟踪推荐系统中的用户兴趣方面是有效的。

图片名称

图3 (a) 在三个工业数据集中,随着模型大小的增加,性能得到提升。(b) 线上学习模型与静态模型之间的性能比较。AUC值越高越好。

系统可扩展性受限于内存和长反馈循环(Long Feedback Loop)。采用在线训练推荐DL模型并将它们的sparse嵌入表扩展到多太字节的第一个挑战是:在内存效率和模型准确性之间进行权衡。如第II-A节所讨论的,sparse特征的sparseID不能在不考虑它们的重要性以及时间访问模式(例如,视频ID可能代表一个只应推广几天的趋势视频)的情况下均匀映射到哈希桶。随着在线训练的进行,sparse嵌入表动态增长,不同嵌入向量的数量增加得更快。这两种效应都导致嵌入表内哈希冲突的可能性高,模型性能下降。现有的开源训练和服务系统的另一个问题是:它们在部署大规模在线训练模型和支持实时数据反馈循环方面的效率低下[25]、[26]。在工业生产环境中,为了实现快速模型恢复和运行A/B测试以衡量它们的模型性能,需要同一任务的多个模型版本共存。因此,Kraken被重新设计,以超越以前的系统,支持具有超过数千亿参数的在线训练和服务推荐模型,同时保持稳定的模型准确性。

三、Kraken设计原则

为了构建一个内存高效的大规模推荐系统的在线学习系统,Kraken中的以下设计原则对于实现不仅系统可扩展性而且高模型质量至关重要:

  • 1) 减少哈希冲突并在特征间共享内存空间。Kraken提出了一种动态地接纳和驱逐sparse嵌入的技术,以避免不必要的哈希冲突,并将所有sparse嵌入存储在全局共享的嵌入表中,以提高内存利用率。通过这种方式,Kraken不为每个sparse特征显式限制哈希桶的大小,而是允许嵌入表在在线学习过程中弹性和自动地调整大小
  • 2) 感知sparse性的训练框架。在拥有超过$10^{11}$的sparse特征的生产推荐系统中,sparse参数占内存资源的99.99%以上。Kraken引入了一个感知sparse性的训练框架来减少训练期间的内存使用。在这个框架下,我们还提出了一种新的优化器rAdaGrad,以进一步减少训练期间与sparse嵌入相关的内存使用。
  • 3) 高效的持续部署和实时服务。在在线学习过程中,训练服务器中的模型始终从新生成的用户数据中学习。Kraken提出了一种高效的方法来部署不断更新的模型,而不会滞后。

为了最小化由哈希冲突导致的准确性损失,我们引入了无冲突但内存高效的嵌入结构——全局共享嵌入表(GSET),如图4所示。

图片名称

图4 Kraken的全局共享嵌入表(Global Shared Embedding Table,简称GSET)

通常增加嵌入表的容量用于减少哈希冲突。然而,很难预测表的大小,因此在部署前设置一个高(通常是恒定的)嵌入表大小,特别是对于在线学习系统来说,这是不高效的。相反,我们提出了GSET,以便:

  • 1)通过全局映射器将键值获取操作与特征嵌入过程解耦,它为每个嵌入表提供了高逻辑容量,通过在它们之间共享内存来灵活地调整物理内存占用,
  • 2)执行自适应条目替换算法,即特殊的特征接纳和驱逐策略,确保在长期执行期间内存占用低于预设阈值。

GSET中的全局映射器:

GSET中的全局映射器将键值获取和嵌入解耦。它以特征的名称和值(例如,‘UID’和‘001’作为特征用户ID)作为输入,并返回一个由预设映射方式生成的格式化键,该键被发送到后端的内存中键值存储系统,该系统负责嵌入管理,而不是传统结构中的普通数组。键的格式主要由两个表示特征名称和特征值的域组成,并且对于模型开发人员的特殊要求具有高度的可配置性。例如,可以分别设置特征值域的宽度,以控制逻辑容量并满足不同特征的倾斜需求。有了全局映射器,GSET允许每个sparse特征的弹性增长和收缩,共享它们之间的内存资源,而不是手工制作不同大小的嵌入表。

特征接纳和驱逐:

为了在长期执行期间控制内存使用,GSET在整个在线学习过程中对所有活动特征执行不同的自适应条目替换算法。自适应条目替换算法利用sparse特征的不同特征来决定如何接纳和驱逐sparse嵌入,包括它们的频率、持续时间、特征重要性等。

  • 例如,许多sparseID在我们的生产数据集中只出现一次,这些不应该被添加到GSET中。
  • 另一个例子是,一些与趋势视频相关的sparseID在这些视频从数据库中退役后就不再需要了。

基于这种领域知识(即特征工程),机器学习工程师为每类sparse特征定制条目替换算法,以最大化模型性能。

GSET中使用的自适应特征接纳政策是:过滤掉频率低的sparseID。由于跟踪永远不会有任何实际用途的稀有特征的统计数据成本很高,GSET支持如[9]中介绍的概率基过滤器。概率基过滤器以概率$ p $接纳不在GSET中的sparseID。sparseID需要被看到的次数遵循几何分布,期望值为$ \frac{1}{p} $。有了这些过滤器,低频ID将在训练过程之前被丢弃,这消除了多余的计算和内存使用。

GSET中使用的传统内存缓存的条目替换算法(如LFU和LRU)旨在最大化缓存命中率,只考虑每个缓存条目被引用的频率。相反,GSET使用在线学习过程中获得的额外信息来确定达到内存限制后条目驱逐的顺序,这称为特征评分方法。GSET为每个sparseID维护一个特征分数,该分数由包含sparseID的训练样本数量以及这些样本的最近程度决定。在间隔期间,GSET通过以下公式更新每个sparseID的特征分数:

\[S^{t+1}_L = (1 - \beta)S^t_L + \beta (c^+ r + c^-)\]

其中:

  • $ \beta $:是时间衰减率,
  • $ r $:是重要性权重,
  • $ c_+, c_- $:分别是在此间隔中包含sparseID L的正例和负例的数量。

当$ \beta = 1 $且$ r = 1 $时,特征评分方法等同于LFU。给正例和负例分配不同权重的原因是正例(例如,点击,点赞)相对较少,在模型预测中更有价值。这与处理不平衡数据集时的下采样技术[9]有相似之处。

在我们的生产工作负载中,我们发现另外两种启发式方法很有用。

第一种称为基于持续时间的政策,它为每个已更新的sparseID设置一个过期时间戳。在使用特征评分方法驱逐sparseID之前,垃圾收集器将首先回收那些已过期的sparseID。这是因为许多sparse特征有明显的生命周期,它们在训练日志中的出现在短时间内迅速消失。机器学习工程师可以在过去的数据分析上运行离线分析,以估计它们的平均持续时间,并为每个sparse特征设置一个最优的持续时间值。例如,我们网站上的许多视频在发布两天后没有视频点击或观看。因此,与视频ID相关的一些特征可以在两天后安全地回收。

第二个优化称为基于优先级的政策,它为有限大小的sparse特征设置驱逐优先级类别。在我们的生产环境中,sparse特征通常被分类为两个优先级类别:高优先级和低优先级。当GSET达到内存限制时,只有低优先级的sparse特征被特征评分方法驱逐。sparse特征的优先级通常由机器学习工程师的领域知识和特征重要性算法决定。在开始在线训练之前,可以使用[27]-[29]中的特征选择现代方法在离线分析中估计特征重要性。然后,可以在它们的总大小不超过某些内存限制的约束下,将顶级特征列表分组到高优先级类别中。例如,在我们的生产工作负载中,关闭与用户相关的特征(例如,用户ID、城市级别)的驱逐有助于提高模型准确性。

C. 高效的持续部署和实时服务

在本节中,我们主要介绍Kraken中为实时服务大规模推荐模型而构建的系统组件。

以前的服务系统设计[25]、[26]在一台预测机器内保持多个模型版本,无法支持需要跨节点共享的大规模推荐模型。

一种简单的方法是同位部署(图5a)(Co-located Deployment),它直接在推理服务器中处理分片模型。具体来说,每个推理服务器维护一个完整的dense部分和一个sparse部分分片。在预测时,它从其他对等方获取所需的参数,并在本地进行预测。然而,这种直接的方式在面对不断更新的模型时引入了相对较高的财务成本。

  • 一方面,每个推理服务器需要高容量的DRAM来存储一部分sparse参数。
  • 另一方面,不断的模型更新影响推理服务器,浪费它们的计算资源和NIC带宽。

为了在提供生产环境所需的功能的同时增强可扩展性,我们重新设计了预测系统,将存储服务和推理服务分别在不同的服务器上处理,称为非同位部署

图片名称

图5 预测系统的两种架构。基线方案直接将所有参数分区到不同的推理服务器上,而Kraken则将sparse嵌入的存储与模型预测的计算解耦。

非同位部署(Non-Colocated Deployment)。如图5b所示,Kraken的预测系统构建为两个服务:预测参数服务器(简称PPS)和推理服务器。

  • 与训练中使用的参数服务器架构类似,预测参数服务器存储分片模型和嵌入表
  • 为了进一步减少请求延迟并节省NIC带宽,推理服务器缓存模型的某些部分,包括整个dense部分和频繁访问的嵌入向量

当接收到包含sparse特征ID列表的请求时,推理服务器从预测参数服务器获取所需的sparse嵌入向量,然后执行模型推理。使用非同位部署的主要好处是允许这两个服务使用不同的硬件资源分别进行扩展。预测参数服务器需要大内存和高网络带宽,而推理服务器主要受计算资源的限制。因此,预测参数服务器可以使用高内存实例,推理服务器可以利用具有高计算能力的机器。

通过非同位部署,我们有两个额外的机会进一步优化服务系统。

  • 一方面,为了增强局部性和负载均衡,Kraken支持按特征放置策略来根据它们的访问模式分布不同类型的参数。这个策略可以将一起访问的参数分组,并将它们放入同一个分片以获得更好的访问局部性。例如,一些用户端的二元sparse特征,如关注列表、收藏列表,通常以结合用户ID和其他项目ID的形式出现。因此,基于用户ID对这些sparse特征进行分片可以增强局部性,因为它们通常对同一用户一起访问。对于极受欢迎的参数,它们甚至可以在PPS的每个分片中复制,以减少热点并实现更好的负载均衡。
  • 另一方面,虽然我们的分布式在线训练系统实现了对机器学习模型的更频繁更新,但也希望以分钟级别的延迟部署新训练的模型进行在线服务。为了同时减少负载并实现实时模型更新,Kraken的训练子系统采用不同的更新策略执行增量模型更新。对于模型的sparse部分,不是每次都传输多太字节嵌入表的完整副本,而是每个嵌入向量的更新将触发一个包含新值的更新消息,然后该消息将被发送到所有下游预测服务器以进行更新。对于模型的dense部分,由于它们的参数比sparse参数的波动性小,整个dense参数的副本将每几秒钟批量更新一次。

IV. 实现

Kraken使用C++11和Python实现。Kraken训练系统的初始版本实现了自己的工作引擎和参数服务器。然而,为了利用TensorFlow生态系统带来的好处,Kraken的新版本被构建为TensorFlow的插件,与TensorFlow的API完全兼容。该插件通过TensorFlow的C++底层API实现为定制操作符,这些操作符与Kraken的参数服务器交互,以执行sparse嵌入向量和dense变量的不同操作。为了预取和批量处理嵌入向量,我们还实现了嵌入缓存来存储在小批量中访问的嵌入向量。类似于Horovod[33],TensorFlow插件在Python API层面添加了钩子和变量代理,以调度工作器发送梯度和参数服务器发送模型参数的时间和通信模式

训练和服务于参数服务器的实现共享相同的代码库。参数服务器的核心(core)是一个高性能的易读键值存储,可以执行梯度聚合算法以及自适应特征管理算法。对于GSET,我们没有像[34]那样构建共享内存池,而是通过直接将所有参数分区到不同服务器来维护一个虚拟表,这是由于键值的简单语义。对于接纳算法,不需要维护任何状态,对于驱逐算法,所有策略都可以通过每个特征仅4字节(用于存储16位时间戳和16位特征分数)来支持。考虑到常见的128或256字节嵌入大小和由感知sparse性训练框架节省的空间,4字节的开销是微不足道的。

除了核心运行时,推理服务器的实现针对生产环境中的推荐任务进行了高度优化,可以支持包括CPU、GPU和FPGA在内的异构计算设备。Kraken的消息系统被构建为一个通用基础设施,不仅支持模型部署,还支持数据分发到需要可扩展解决方案的其他存储系统(例如,存储项目和用户特征的索引服务和用户档案服务)。

Kraken支持在线训练算法的异步和同步模式。在生产中,随着我们的模型扩展并需要更多的工作器,我们发现异步模式具有更高的训练速度,并且对机器故障更加健壮。因此,异步在线训练成为了训练我们的推荐模型的默认选项。

V. 评估

我们首先评估我们提出的技术的好处,然后报告Kraken在生产中实际应用的性能。

A. 实验设置

评估平台。我们在拥有64台服务器的集群上评估Kraken,除了从生产系统直接收集指标的生产性能评估外。集群中的所有服务器都配备了512GB DRAM和两个2.5GHz Intel(R) Xeon(R) Gold 6248 CPU,每个CPU有20个核心。服务器使用10 Gbps以太网连接。如果没有特别说明,每台服务器都配备了四个训练工作进程和一个参数服务器进程。

数据集。我们在公共数据集和生产数据集上衡量Kraken的性能,以便我们的实验可以轻松复现,同时展示实际工业场景中的性能。使用的公共数据集包括Criteo广告数据集、Avazu CTR数据集和MovieLens25M。Criteo广告数据集[35]在评估推荐模型时非常流行,并且很快将作为标准基准包含在MLPerf基准[36]中。Avazu CTR数据集[37]包含了一个领先广告平台的站点、应用和设备上的点击数据,共有11天。MovieLens-25M[38]通常作为一个稳定的基准数据集,包含2500万个电影评分,评分范围从1到5,以0.5为增量。在这里,我们将高于3的样本标记为正样本,其余为负样本,并将其训练为二元分类模型。两个生产数据集是从两个独立的现实世界推荐服务中收集的:Explore Feed和Follow Feed。这两项服务都向用户推荐与视频相关的内容。表III总结了不同数据集的特点。请注意,Explore Feed数据集需要5亿参数来构建一个合理的推荐模型,而Follow Feed数据集需要500亿参数(多100倍)。我们应用常见的指标,AUC和Group AUC[39](GAUC),来评估模型的准确性。GAUC是所有用户AUC的加权平均值,以每个用户的样本数量为权重。因此,GAUC比AUC更能指示现实世界推荐系统中的模型性能。所有数据集都以在线学习的方式学习,即每个样本只训练一次。

实验比较了四种工业模型,DNN、Wide & Deep[16]、DeepFM[17]、DCN[40]在不同数据集上与Kraken和TensorFlow的结果。如果没有特别说明,默认的微基准模型是DeepFM。其他模型有类似的结论,因此我们省略它们以节省空间。更详细的设置,如模型的超参数,可以在我们的Artifact Description中找到,以便复现。

B. 端到端系统性能

我们评估Kraken的好处,包括系统方面和模型方面的表现。我们分别在TensorFlow和Kraken中应用Adam和混合优化器。具体来说,在Kraken中,dense部分的优化器是Adam,而sparse部分是rAdaGrad。选择Adam的原因是它是最受欢迎的优化器,并且在许多研究论文[17]、[41]中因其出色的收敛速度和少调整而被使用。更多与优化器相关的评估,如AdaGrad或SGD,将在后面的章节中介绍。为了公平比较,TensorFlow的嵌入表大小被设置为使内存消耗接近Kraken(可以容纳60%的所有原始特征)。

模型准确性。我们比较Kraken和TensorFlow以评估模型的准确性。由于计算资源有限,我们不对Follow Feed进行评估,因为它需要64台服务器从头开始训练,并持续一个月以验证一个可能的配置,使用一个50B参数模型。对于TensorFlow,我们尝试了五种不同的嵌入表大小组合,并仅显示最佳结果。需要强调的是,在有限内存的情况下,调整TensorFlow的嵌入表大小以获得好的模型是费力且耗时的。

图6显示了Kraken与TensorFlow相比在不同模型上带来的模型准确性提升。对于评估的四个工作负载,Kraken在公共数据集上比TensorFlow高出0.46%到6.01%的AUC,在Explore Feed上比TensorFlow高出1.64%到2.01%的GAUC,这是一个巨大的提升(在生产中,超过0.5%的提升是显著的)。Kraken的提升主要来自GSET的无哈希冲突设计和感知sparse性训练框架的适应性。此外,它显著减少了时间和计算资源,允许每个嵌入表的弹性增长,消除了对嵌入表大小的广泛调整。图6显示GSET可以实现比最佳人工调整的嵌入表更好的模型性能。

图片名称

图6

系统开销。然后我们评估GSET对底层TensorFlow施加的开销,这可能影响端到端的训练速度。图7a通过在Explore Feed上变化训练服务器的数量,比较了Kraken和原生TensorFlow的吞吐量(即每秒样本数)。如图所示,Kraken的吞吐量始终接近或优于TensorFlow,从而洞察到Kraken的非常小的额外开销。此外,随着工作器数量的增加,Kraken保持线性增长,而当工作器数量增长到大约75时,TensorFlow的增长率下降。

图片名称

图7

可扩展性。图7b显示Kraken在Follow Feed数据集上线性扩展。不幸的是,TensorFlow由于内存不足,不支持如此大模型的训练,当需要将如此多的特征列适应到原生TensorFlow嵌入表中时。Kraken显示出对大规模模型的更好适应性。

C. GSET评估

在本节中,我们评估GSET的设计。为了消除感知sparse性优化器的影响,我们在Kraken和TensorFlow中应用了相同的Adam优化器。如果没有特别说明,Kraken和TensorFlow使用相同的内存(足以存储60%的所有特征)。

GSET的内存效率。为了分析GSET在在线学习中的内存效率,我们在不同内存占用下(即,最多保存所有原始ID的几个比例)比较了使用Kraken和TensorFlow在Criteo数据集上不同模型的AUC。对于Kraken,特征接纳概率被设置为1,并启用了驱逐机制。如图所示,在不同的内存占用下,Kraken始终优于TensorFlow超过0.61%甚至高达3.72%。通常情况下,内存越少,GSET的提升越多。这是因为在原生TensorFlow中更激烈的哈希冲突使得模型更难学习代表输入特征的良好嵌入条目。GSET的设计减少了嵌入的哈希冲突,并在在线学习过程中学习了一组有效的特征。图8还说明了GSET的弹性增长设计可能是调整嵌入表大小的完美解决方案。

图片名称

图8

特征接纳的效果。图9展示了在Explore Feed数据集中不同特征接纳概率P的效果。从图9a中,模型性能似乎不受接纳概率的影响。这是合理的,因为这些低频特征很少对整个模型做出贡献。Kraken简单地丢弃这些特征以避免频繁的过期。图9b显示了在最后一个训练小时中出现的不同频率级别的特征数量。有趣的是,不同接纳概率的低频特征数量几乎没有差异,这可能最初看起来违反直觉,但回想起来,尽管一些低频特征被过滤掉,一些相对高频的特征也更少地进入系统,从而成为新的低频特征。

图片名称

图9

特征驱逐的效果。我们还进行了因素分析,以了解在保持内存使用的同时,每种特征驱逐策略对模型性能的贡献。在这个实验中,我们省略了Criteo Ad数据集,因为它的训练样本是特征匿名的,不包含时间戳信息,这将使特征驱逐盲目。

我们通过测量以下三种设置的模型性能差距来分解基线GSET仅具有LFU过期策略和结合所有三种驱逐策略的最佳性能:

  • 特征评分策略(F)进一步考虑了正负样本的不同优先级。重要性权重r是从验证数据集中的[1, 3, 5]中选择的。特征分数每天衰减10%。
  • 基于持续时间的策略(D)为每个特征类别设置不同的过期持续时间。我们预先采样10%的数据,分析具有相同ID的相邻样本之间的间隔时间分布,并取99.9百分位作为该特征类别的过期时间。
  • 基于优先级的策略(P)禁止消除在Avazu和Explore Feed数据集中占用总内存不到10%的用户相关特征,基于我们的实践。然而,我们只在MovieLens数据集中禁用了占用较少内存并且访问更频繁的物品ID的驱逐。这是因为用户相关特征占据了高达50%,禁止驱逐它们将占用其他特征的内存,导致不满意的准确性。

图10表明,我们的特征评分策略在不同数据集上一致优于LFU策略,三种不同策略的累积可以结合它们的优势。可以得出结论,ML算法的领域知识和数据分布的感知可以使驱逐更智能,这就是为什么我们需要一个可配置的特征驱逐组件,以便ML工程师灵活定制。请注意,策略的超参数既与模型有关,也与数据集有关。我们在这里为了简单起见进行了经验设置,并且可以使用优化系统[42]进行更精细的调整。

图片名称

图10

D. 感知sparse性训练框架的效果

接下来,我们展示Kraken的感知sparse性训练框架能够像原始优化器一样正确地收敛模型,并且在更少的内存资源下提供更好的准确性。在这个实验中,我们只关注优化器,因此我们为GSET提供了足够的内存,并关闭了特征接纳和驱逐功能。

表IV显示了在三个公共数据集上不同原始优化器和感知sparse性优化器的模型性能和内存消耗。在相同的内存消耗下,我们提出的混合优化器总是能够取得更好的性能,除了在两个负样本上有一些相似的性能。我们提出的组合Dense(Adam)&Sparse(rAdaGrad)在Test AUC和内存消耗方面都优于其他所有基线优化器和混合优化器。

它能够在减少3倍内存使用的同时,提供与Adam(正常优化器中最好的)一样高模型性能。尽管组合Dense(Adam)&Sparse(SGD)的OSPs稍少,但由于缺乏适应性,最终模型性能更差。它未能在实时推荐场景的sparse数据中很好地学习。Kraken的rAdaGrad提供了最小的存储开销和学习的适应性。

感知sparse性训练框架在性能和内存效率上的提升与第三节B部分一致。结果表明,Kraken的算法不依赖于模型或数据集。此外,感知sparse性训练框架节省的内存资源可以用来替换更多的sparse参数,从而提升模型性能。

E. 大规模在线预测评估

在本节中,我们构建了一个成本模型来评估Kraken预测系统的两种部署策略。我们考虑了两个集群组件的成本:预测参数服务器和推理服务器,并计算了云中每美元的预测吞吐量。基线是一个运行同位部署策略的预测系统。Kraken的非同位部署使其能够灵活地为预测参数服务器和推理服务器配置不同的服务器。这很重要,因为预测参数服务器是内存受限的,需要大内存,而推理服务器是计算受限的,不需要大内存。然而,对于基线的同位部署,所有推理服务器都需要是具有大内存的计算dense型服务器。 我们以一个16分片的Follow Feed模型为例进行进一步的成本建模。基于CPU和NIC带宽利用率的计算,一组16个预测参数服务器可以承载的推理服务器的最大数量估计为384。表V总结了使用两种不同部署策略的硬件成本。非同位部署在性价比上优于同位部署1.3倍(使用AWS价格数据[43])或2.1倍(使用阿里云价格数据[44])。从上述数据和分析中,我们得出结论,非同位部署在大规模推理集群中更有效,实现了更低的硬件成本。显然,通过将推理服务器从不断更新参数中解耦出来,Kraken实现了低成本和出色的推理性能。

F. 生产评估

由于Kraken已经在生产中部署了两年,我们报告了在几个现实世界应用中使用Kraken的性能指标,以展示其成功。

1) 在线A/B测试结果:我们选择了三个由Kraken支持的代表性应用:视频分享、社交网络和游戏平台。表VI显示了在使用Kraken后它们的关键业务指标的增长,如下所述:

  • 视频分享是一个相关视频推荐应用,在使用户观看共享视频后提供更多视频建议。其关键业务指标是每个视频的平均播放次数(平均视频播放次数)。Kraken实现了视频播放次数51%的增长,并显著提高了用户参与度。
  • 社交网络是在我们的平台上向用户推荐潜在社交联系的服务。新社交联系的平均数量是评估此服务的关键指标。Kraken将核心指标提高了1.35%,使更多用户连接到其他用户。(1.35%是显著的,与旧系统中通常的0.1%改进相比。)
  • 游戏平台是一个托管不同数字游戏的在线平台,Kraken用于在其Feed中生成个性化的游戏视频推荐。其关键指标是用户在阅读Feed上花费的总时间(Feed上总时间花费)。Kraken在关键指标上提高了24.41%,显示出提高用户粘性的有效性。

2) 日常监控结果:我们还通过监控Kraken服务的推荐模型的准确性以及其在整个一天的服务吞吐量和延迟来报告生产中的Follow Feed应用的性能。(Follow Feed应用与第V-B节中的离线评估中使用的应用相同。)

  • 模型准确性:图11a显示了Kraken生成的平均预测点击率(CTR-P)和Follow Feed中项目的平均点击率真实值(CTR-GT)。高点击率通常意味着高用户参与度,更准确的CTR预测有助于项目推荐。如图所示,CTR-P曲线与CTR-GT曲线非常接近,表明Kraken的模型预测准确性高。
  • 系统性能:图11b展示了Kraken的系统吞吐量(即推理请求的数量)以及平均延迟和尾部延迟(P99)随时间的变化。在这一天中有两个明显的高峰,12:00至14:00和20:00至23:00。在后一个时期(图11b中的阴影区域),称为高峰时段,吞吐量高达40k QPS(每秒查询次数),是平均吞吐量的两倍。与此同时,尽管吞吐量急剧上升,Kraken仍然很好地控制了平均延迟和尾部延迟(P99)。

图片名称

图11

VI. 相关工作

尽管系统和架构社区已经投入了大量精力对用于计算机视觉或自然语言处理的深度学习进行性能分析和优化,但相对较少的关注集中在在线学习和实时推荐系统中大规模深度学习模型的实时服务上。最相关的工作是Facebook的DLRM[1]、[13],它指出了在现代生产规模下训练推荐系统的DNN所面临的独特挑战。它提供了详细的性能分析,表明推荐DL模型需要更大的存储容量并产生不规则的内存访问。他们使用蝴蝶洗牌操作符在嵌入表上实现模型并行以缓解内存限制。XDL[14]和Parallax[45]明确区分DL模型中的dense部分和sparse部分,并尝试通过许多sparse变量改进训练模型。Parallax通过使用AllReduce架构处理dense变量和PS架构处理sparse变量,为NLP的同步训练采用混合架构。然而,上述系统大多关注批量训练模型,其嵌入表大小是固定的,没有动态增长和全局空间共享。这些系统错过了实时训练和成本有效地服务10-TB DL模型的机会。传统的深度学习框架如[11]、[12]、[46]、[47]被全球的机器学习科学家和工程师使用。然而,它们在面对推荐系统中大规模实时推荐挑战时都无法扩展。HugeCTR[48]是NVIDIA为推荐系统设计的高效率GPU框架。HugeCTR将整个嵌入表分布到多个GPU的高带宽内存(HBM)中,以加速CTR模型的训练。然而,HugeCTR受到HBM大小的限制,因为所有sparse参数都必须在HBM内维护。在这种情况下,Kraken提供的内存高效学习可以被完美应用,并在节省计算资源开销方面发挥关键作用。

参数服务器(PS)[21]作为数据并行分布式训练领域中的代表性架构之一,在Kraken中作为底层通信模型使用。多年来,有许多先前的工作在许多方面优化了并行训练,包括利用集群中的GPU[24]、[49]、网络[50]和调度[51]–[53]。这些工作提出的技术与我们的工作正交,可以用来进一步改进Kraken中的训练。随着持久内存(PM)如英特尔Optane DC持久内存的出现,将传统的基于PM的键值存储[54]适应到Kraken并将是一件有趣的事情,并在新硬件带来的高能力的基础上缓解内存短缺。

持续学习已被证明是跟上推荐系统中用户兴趣快速变化的有效解决方案[6]、[7]、[9]。Continuum[55]是一个通用的持续学习系统,通过封装不同的学习框架,使用批量模式在新数据集上重新训练模型。与Kraken相比,Continuum更新模型的频率较低,并且不适用于大规模推荐模型。

VII. 结论

通过利用系统和训练算法的共同设计,Kraken提供了一个端到端的解决方案,用于实时训练和大规模推荐模型的服务。Kraken的训练系统实现了一个参数服务器,允许sparse嵌入表动态增长,并运行自动特征选择,在持续训练期间保持合理的内存大小。我们还提出了一个感知sparse性的框架,利用推荐模型的属性,在训练期间减少数据传输和内存占用。此外,Kraken的预测系统支持及时有效地部署大规模模型。Kraken已成功部署在生产中,用于广泛的推荐任务,并被证明在这些任务中迭代和服务大规模DL模型方面非常有效。

参考

介绍

google的Jeffrey Dean在2012《Large Scale Distributed Deep Networks》这篇经典paper中提出了大规模分布式DNN的设计。我们回顾一下:

摘要

最近在无监督特征学习和深度学习领域的工作表明,能够训练大型模型可以显著提高性能。在本文中,我们考虑了使用数万个CPU核心训练具有数十亿参数的深度网络的问题。我们开发了一个名为DistBelief的软件框架,可以利用数千台机器的计算集群来训练大型模型。在此框架内,我们开发了两种大规模分布式训练算法:

  • (i)Downpour SGD,一种支持大量模型副本的异步随机梯度下降过程;
  • (ii)Sandblaster,一个支持多种分布式批量优化过程的框架,包括L-BFGS的分布式实现。

Downpour SGD和Sandblaster L-BFGS都增加了深度网络训练的规模和速度。我们已经成功地使用我们的系统训练了一个比文献中先前报道的深度网络大30倍的网络,并在ImageNet上实现了最先进的性能,ImageNet是一个包含1600万张图像和21k类别的视觉对象识别任务。我们展示了这些相同的技术如何显著加速了一个更适度规模的深度网络的训练,用于商业语音识别服务。尽管我们专注于这些方法在训练大型神经网络方面的性能,但底层算法适用于任何基于梯度的机器学习算法。

1 引言

深度学习和无监督特征学习在许多实际应用中显示出巨大的潜力。在几个领域已经报告了最先进的性能,从语音识别[1, 2]、视觉对象识别[3, 4]到文本处理[5, 6]。还观察到,增加深度学习的规模,无论是训练样本的数量、模型参数的数量,还是两者都增加,都可以大幅提高最终分类精度[3, 4, 7]。这些结果导致了对扩大这些模型的训练和推理算法的兴趣激增[8],并改进适用的优化程序[7, 9]。GPU的使用[1, 2, 3, 8]是近年来的一个重大进步,使得训练适度规模的深度网络变得实际。GPU方法的一个已知限制是,当模型不适合GPU内存时(通常小于6GB),训练加速很小。为了有效地使用GPU,研究人员经常减小数据或参数的大小,以便CPU到GPU的传输不是一个显著的瓶颈。虽然数据和参数的减小对于小问题(例如语音识别中的声学建模)效果很好,但对于具有大量样本和维度的问题(例如高分辨率图像)则不太吸引人。

在本文中,我们描述了一种替代方法:使用大型机器集群来分布式训练和推理深度网络。我们开发了一个名为DistBelief的软件框架,它支持在机器内部(通过多线程)和机器之间(通过消息传递)的模型并行性,框架管理并行性、同步和通信的细节。除了支持模型并行性(s model parallelism),DistBelief框架还支持数据并行性(data parallelism),其中使用多个模型副本来优化单个目标。在此框架内,我们设计并实现了两种新的大规模分布式训练方法:

  • (i)Downpour SGD:一种异步SGD过程,利用自适应学习率并支持大量模型副本;
  • (ii)Sandblaster L-BFGS:一种使用数据和模型并行性的L-BFGS的分布式实现。

Downpour SGD和Sandblaster L-BFGS与更传统的SGD和L-BFGS实现相比,都获得了显著的速度提升。我们的实验揭示了关于大规模非凸优化的几个令人惊讶的结果。

  • 首先,异步SGD很少应用于非凸问题,但对于训练深度网络非常有效,特别是当与Adagrad[10]自适应学习率结合时。
  • 其次,我们展示了:如果有足够的资源,L-BFGS与许多SGD变体相比具有竞争力或更快。

关于深度学习中的具体应用,我们报告了两个主要发现:我们的分布式优化方法既可以大大加速适度规模模型的训练,也可以训练比以往更大的模型

  • 为了说明第一点,我们展示了我们可以使用机器集群以不到GPU所需时间的1/10来训练一个适度规模的语音模型,达到相同的分类精度。
  • 为了说明第二点,我们训练了一个超过10亿参数的大型神经网络,并使用这个网络在ImageNet数据集上大幅提高了最先进的性能,这是计算机视觉中最大的数据集之一。

2 先前的工作

近年来,商业和学术机器学习数据集以前所未有的速度增长。作为回应,许多作者已经探索了扩大机器学习算法以应对这些数据的洪流[11, 12, 13, 14, 15, 16, 17]。这些研究中的大部分集中在线性、凸模型[11, 12, 17]上。在凸情况下,分布式梯度计算[18]是自然的第一步,但有时由于同步问题而遭受减速。有一些有希望的努力来解决这个问题,例如在异步随机梯度下降中的无锁参数更新,例如Hogwild[19]。

不幸的是,将这些方法扩展到密集的非凸问题,例如在训练深度架构时遇到的问题,基本上是未知的领域。特别是,不知道是否可能在多个局部最小值存在的情况下平均参数或执行密集的异步参数更新。在深度学习的背景下,大部分工作集中在在单台机器上训练相对较小的模型(例如,Theano[20])。扩大深度学习的一个有趣的建议是使用GPU农场来训练许多小型模型,然后平均它们的预测[21],或者修改标准深度网络使其更易于并行化[22]。

与以前的工作相比,我们的重点是在不限制模型形式的情况下,扩大具有数十亿参数的非常大的模型的深度学习技术。在这种情况下,模型并行性,类似于[23]的精神,是一个基本成分,但必须与巧妙的分布式优化技术相结合,这些技术利用数据并行性

我们考虑了许多现有的大规模计算工具,将其应用于我们的问题,MapReduce[24]和GraphLab[25]是显著的例子。我们得出结论:

  • MapReduce,旨在并行数据处理,不适合深度网络训练中固有的迭代计算;
  • 而GraphLab,旨在通用(非结构化)图计算,不会利用在深度网络中通常发现的结构化图中的计算效率。

3 模型并行性

为了促进非常大的深度网络的训练,我们开发了一个名为DistBelief的软件框架,它支持在神经网络和分层图形模型中的分布式计算。用户定义在模型的每个层的每个节点处进行的计算,以及在计算的向上和向下阶段应该传递的消息。对于大型模型,用户可以将模型分割到几台机器上(图1),以便将不同节点的计算责任分配给不同的机器。框架自动使用所有可用核心在每台机器上并行化计算,并在训练和推理期间管理机器之间的通信、同步和数据传输。

图片名称

图1 DistBelief中模型并行性的例子。这里展示了一个具有局部连接性的五层深度神经网络,它被分割在四台机器(蓝色矩形)上。只有那些边跨越分区边界(粗线)的节点才需要在机器之间传输它们的状态。即使在节点有多个边跨越分区边界的情况下,其状态也只发送给边界另一侧的机器一次。在每个分区内部,单个节点的计算将在所有可用的CPU核心上并行化。

在多台机器上分布深度网络的性能优势取决于模型的连接结构和计算需求。具有大量参数或高计算需求的模型通常从访问更多的CPU和内存中受益,直到通信成本占主导地位。我们已经成功地在DistBelief框架中运行了多达144个分区的大型模型,并取得了显著的速度提升,而更适度规模的模型在多达8或16个分区上显示出不错的速度提升。(见第5节,在“模型并行性基准测试”标题下查看实验结果。)显然,具有局部连接结构的模型比全连接结构更容易进行广泛的分布,因为它们的通信需求较低。提速不理想的典型原因是:不同机器之间的处理时间有差异,导致许多机器需要等待那台最慢的机器,以便完成给定阶段的计算。尽管如此,对于我们最大的模型,我们可以有效地使用32台机器,每台机器平均使用16个核心,总共512个CPU核心训练一个大型神经网络。当与下一节描述的分布式优化算法结合使用时,这些算法利用整个神经网络的多个副本,可以利用数万个CPU核心来训练单个模型,从而显著减少整体训练时间。

4 分布式优化算法

在DistBelief框架内并行化计算使我们能够实例化和运行比以前报告的更大的神经网络。但是,为了在合理的时间内训练如此大的模型,我们不仅需要在单个模型实例内进行并行计算,也需要跨多个模型实例进行分布式训练。在本节中,我们描述了这第二级别的并行性,我们使用一组DistBelief模型实例或副本来同时解决单个优化问题。

我们比较了两种大规模分布式优化程序:

  • Downpour SGD,一种在线方法
  • Sandblaster L-BFGS,一种批量方法

两种方法都利用了集中式分片参数服务器的概念,模型副本使用它来共享它们的参数。两种方法都利用了DistBelief在每个单独副本内允许的分布式计算。但最重要的是,两种方法都被设计为容忍不同模型副本的处理速度差异,甚至是模型副本的大规模故障,这些副本可能被离线或随机重启。

从某种意义上说,这两种优化算法实现了数据并行性的智能版本。两种方法都允许我们同时在许多模型副本中的每一个中处理不同的训练样本,并定期组合它们的结果来优化我们的目标函数。

4.1 Downpour SGD

随机梯度下降(SGD)或许是训练深度神经网络最常用的优化过程[26, 27, 3]。不幸的是,SGD的传统公式本质上是顺序的,使得它不适用于非常大的数据集,因为在完全顺序化的方式下完成全部数据计算所需的时间是令人望而却步的。

为了将SGD应用于大型数据集,我们引入了Downpour SGD,这是一种使用单个DistBelief模型的多个副本的异步随机梯度下降变体。基本方法如下:我们将训练数据分成多个子集,并在这些子集上运行模型的副本。模型通过一个集中的参数服务器通信更新,该服务器保存模型的所有参数的当前状态,这些参数分布在许多机器上(例如,如果我们有10个参数服务器分片,每个分片负责存储和应用模型参数的1/10的更新)(图2)。这种方法在两个不同的方面是异步的:模型副本独立运行,参数服务器分片也独立运行。

图片名称

图2 左边:Downpour SGD。模型副本异步地从参数服务器获取参数 $ \mathbf{w} $ 并推送梯度 $ \Delta \mathbf{w} $。右边:Sandblaster L-BFGS。一个单独的“协调器”向副本和参数服务器发送小消息,以协调批量优化。

在最简单的实现中,在处理每个小批量之前,模型副本会向参数服务器服务请求其模型参数的更新副本。由于DistBelief模型本身分布在多台机器上,因此每台机器只需要与持有与其分区相关的模型参数的参数服务器分片通信。在接收到其参数的更新副本后,DistBelief模型副本处理一小批数据以计算参数梯度,并将梯度发送到参数服务器,然后参数服务器将梯度应用于模型参数的当前值。

可以减少Downpour SGD的通信开销,通过限制每个模型副本仅在每$ n_{\text{fetch}}$ 步请求更新参数,并且仅在每$ n_{\text{push}}$步发送更新的梯度值(其中:$ n_{\text{fetch}} \text{可能不等于} n_{\text{push}}$。

实际上,获取参数(fetch)、推送梯度(push)和处理训练数据的过程可以在三个仅弱同步的线程中进行(见附录中的伪代码)。在下面报告的实验中,我们为了简单和便于与传统SGD比较,将 $ n_{\text{fetch}} = n_{\text{push}} = 1 $ 固定。

Downpour SGD比标准(同步)SGD更能够抵抗机器故障。对于同步SGD,如果一台机器失败,整个训练过程将被延迟;而对于异步SGD,如果模型副本中的一台机器失败,其他模型副本继续处理它们的训练数据并通过参数服务器更新模型参数。另一方面,Downpour SGD中的多种异步处理引入了优化过程中的大量额外随机性。最明显的是,模型副本几乎肯定是基于一组稍微过时的参数计算其梯度,因为其他模型副本可能在此期间已经更新了参数服务器上的参数。但是,除此之外还有几个其他的随机性来源:由于参数服务器分片独立运行,不能保证在任何给定时刻,每个参数服务器分片上的参数经历了相同数量的更新,或者更新以相同的顺序应用。此外,由于模型副本被允许在单独的线程中获取参数和推送梯度,可能还存在参数的时间戳的额外微妙不一致性。对于非凸问题,这些操作的安全性几乎没有理论基础,但在实践中我们发现放宽一致性要求非常有效

我们发现一种可以大大增加Downpour SGD鲁棒性的技术是使用Adagrad[10]自适应学习率过程。Adagrad不是在参数服务器上使用单一固定学习率(图2中的η),而是为每个参数使用单独的自适应学习率。设:

  • $ \eta_{i,K} $ 为第 $ i $ 个参数在第 $ K $ 次迭代的学习率
  • $ \Delta w_{i,K} $ 为其梯度

则我们设置:

\[\eta_{i,K} = \frac{\gamma}{\sqrt{\sum_{j=1}^{K} (\Delta w_{i,j})^2}}\]

由于这些学习率仅从每个参数的累积平方梯度计算,Adagrad可以很容易地在每个参数服务器分片内局部实现。γ的值,所有学习率的常数缩放因子,通常比没有使用Adagrad时使用的最佳固定学习率大(可能大一个数量级)。Adagrad的使用扩展了可以同时有效工作的模型副本的最大数量,并且结合了仅使用单个模型副本“预热”模型训练的做法,然后释放其他副本,它几乎消除了使用Downpour SGD训练深度网络时的稳定性问题(见第5节的结果)。

4.2 Sandblaster L-BFGS

批量方法已被证明在训练小型深度网络方面表现良好[7]。为了将这些方法应用于大型模型和大型数据集,我们引入了Sandblaster批量优化框架,并讨论了在此框架中使用L-BFGS的实现。

Sandblaster的一个关键思想是分布式参数存储和操作。优化算法(例如L-BFGS)的核心位于协调器进程(图2),它没有直接访问模型参数。相反,协调器发出来自一小组操作(例如,点积、缩放、系数加法、乘法)的命令,每个参数服务器分片可以独立执行这些操作,并将结果存储在同一分片上。附加信息,例如L-BFGS的历史缓存,也存储在计算它的参数服务器分片上。这允许运行大型模型(数十亿参数)而不会因将所有参数和梯度发送到单个中央服务器而产生开销(见附录中的伪代码)。

在典型的L-BFGS并行化实现中,数据被分发到多台机器上,每台机器负责计算特定子集数据示例上的梯度。梯度被发送回中央服务器(或通过树形结构聚合[16])。许多这样的方法等待最慢的机器,因此不适用于大型共享集群。为了解决这个问题,我们采用了以下负载平衡方案:协调器为N个模型副本中的每一个分配一小部分工作,远小于批次总大小的1/N,并且每当副本空闲时就分配新的部分。通过这种方法,较快的模型副本比慢的副本做更多的工作。为了进一步管理批次末尾的慢模型副本,协调器调度多份未完成的部分,并使用最先完成的模型副本的结果。这种方案类似于MapReduce框架中“备份任务”的使用[24]。数据预取,以及通过将数据的连续部分分配给同一工作器来支持数据亲和性,使得数据访问不成问题。与Downpour SGD相比,后者需要相对高频率、高带宽的参数同步与参数服务器,Sandblaster工作器仅在每个批次开始时(当它们被协调器更新时)获取参数,并且仅在完成几个部分后发送梯度(以防止副本故障和重启)。

实验

参考

介绍

ali在2020《XDL: An Industrial Deep Learning Framework for High-dimensional Sparse Data》提出了它们的XDL平台。我们看下它的实现:

摘要

随着数据和计算能力的快速增长,基于深度学习的方法已成为许多人工智能问题的主流解决方案,如图像分类、语音识别和计算机视觉。包括Tensorflow、MxNet和PyTorch在内的几个优秀的深度学习(DL)框架已经开源,进一步加速了社区的进步。然而,现有的DL框架并不针对涉及高维稀疏数据的应用程序设计,这种数据在许多成功的在线业务中广泛存在,如搜索引擎、推荐系统和在线广告。在这些工业场景中,深度模型通常在大规模数据集上进行训练,这些数据集包含高达数十亿的稀疏特征和数千亿的样本,给DL框架带来了巨大挑战。

在本文中,我们介绍了一个高性能、大规模且分布式的DL框架XDL,它提供了一个优雅的解决方案,填补了现有DL框架的通用设计与高维稀疏数据引起的工业需求之间的差距。自2016年以来,XDL已在阿里巴巴成功部署,服务于许多生产环境,如在线广告和推荐系统。在数百个GPU卡上并行运行,XDL可以在仅几个小时内训练具有数千亿参数的深度模型。除了其卓越的性能和灵活性,XDL对开发人员也非常友好。阿里巴巴的算法科学家可以用几行简单的代码开发和部署新的深度模型。XDL API和参考实现作为开源包在2018年12月发布,遵循Apache 2.0许可,并可在https://github.com/alibaba/xdeeplearning找到。

1 引言

在过去的十年中,作为人工智能最令人兴奋和强大的分支之一,深度学习在语音识别、计算机视觉、自然语言处理和医学诊断等领域取得了重要的突破。得益于许多现实世界应用中产生的天文数量级的数据,以及图形处理单元(GPU)等基础设施带来的前所未有的计算能力,以及TensorFlow[1]、MxNet[6]、Caffe[18]等复杂的开源深度学习框架,深度学习技术已被广泛用于解决现实世界问题。

尽管现有的深度学习框架在许多领域取得了巨大成功,但它们并不针对涉及高维稀疏数据的应用程序设计友好。高维稀疏数据在许多互联网规模的应用中广泛存在,如搜索引擎、推荐系统和在线广告。例如,在我们的展示广告系统中,我们每天生成的用户行为日志数据量达到PB。从这些数据中提取的训练样本包含数十亿个特征,而每个样本只有少数几个维度是非零的

与语音识别、计算机视觉和自然语言处理等数据密集型应用不同,这些互联网规模的在线业务,数据稀疏,对深度学习框架提出了独特的挑战。

  • 首先,问题的规模如此之大,以至于无法在单个节点上解决。需要并行化来解决这种规模的问题。
  • 其次,数据极其稀疏。如果处理不当,稀疏性可能导致效率低下。
  • 第三,这些数据的表述并不规范。样本之间可能有重复的特征,给带宽带来高压。

此外,在快速模型演化中,需要添加或删除特征以测试其与目标的相关性。现有的深度学习框架并不设计友好地解决这些挑战。这种差距使得这些业务难以充分从蓬勃发展的深度学习技术中受益。

通过抽象我们在构建在线广告系统的深度学习解决方案中的实践经验,涉及高维稀疏数据的深度模型通常包括两部分:sparse特征学习和dense模型学习。稀疏特征学习是将大量高维稀疏数据转换为密集特征的过程,而密集模型学习则负责基于密集特征学习最优模型结构。这两部分还需要智能连接,以便这两个学习任务可以相互交互以实现最佳结果。

在本文中,我们介绍了XDL,这是一个为涉及高维稀疏数据的学习任务设计的高性能、大规模且分布式的深度学习框架。在稀疏特征学习部分,XDL提供了一个精心设计的分布式系统,对I/O、数据管道、通信和GPU进行了深入优化,提供了极高的效率和可扩展性。用户可以使用embedding字典或像CNN/RNN这样的深度模型将稀疏项目特征或图像/文本特征映射到密集表示。在密集模型学习部分,XDL采用任何开源深度学习框架作为其后端,并采用一种名为桥接技术的新技术支持。有了这种设计,XDL能够加速互联网规模学习问题的培训,并在端到端生产系统中提供服务。

在对XDL进行真实世界数据集评估时,我们发现它至少可以比Tensorflow/MxNet的原生分布式版本快5倍。

自2016年以来,XDL已在我们公司的核心业务中部署,如电子商务推荐和在线广告。在数百台服务器上运行,XDL在仅几个小时内训练了数十亿参数的多个模型。除了其性能和灵活性,XDL致力于隐藏复杂的工程细节。使用XDL,我们的算法科学家可以用几行简单的代码开发和部署新模型。这种系统设计带来了灵活性,并解锁了许多算法创新。

2 相关工作

随着问题规模和复杂性的增加,我们面临着深度学习框架性能和可扩展性的更严重挑战。传统的解决方案如TensorFlow[1]、MxNet[6]和Caffe[18]通常在单机上运行,并为分布式训练提供有限支持。即使有多个GPU,这些单机解决方案也难以处理具有$10^{11}$样本和$10^{10}$参数的模型训练。因此,分布式训练平台似乎是一个有希望的解决方案,许多学者在这方面做出了巨大贡献。

MPI[14]定义了一组接口,用于机器之间的通信和协调。由于其兼容性和可扩展性,MPI在分布式内存应用的并行计算中被广泛使用。许多机器学习框架都建立在MPI之上。

  • Chen等人提出了RABIT[5],这是一个AllReduce库,通过增加容错属性改进了OpenMPI。
  • Adam Coates等人提出了COTS HPC技术[3, 9]。他们在InfiniBand互连上用MPI构建了一个GPU服务器集群。仅用16台机器,这个系统就可以有效地训练具有110亿参数的模型。

然而,这些框架的数据处理部分过于简化。因此,它们既不能充分利用数据特性,也不能处理大规模稀疏数据的工作。此外,基于MPI的系统在容错方面支持不足,因此在一些工业场景中导致了可用性问题。

另一个广泛使用的范式是参数服务器(PS),它可以利用键值存储适当处理稀疏数据。

  • Dean等人[11]首次引入了PS,并使用downpour SGD在谷歌训练大规模深度网络。
  • Petuum[30]引入了有界延迟模型到PS架构[23],实现了异步计算。

参数服务器在工业中也广泛使用。

  • DSSTNE[2]是亚马逊设计的DL框架,特别适用于大型稀疏数据集。它提供了极其高效的自动模型并行多GPU支持和100%确定性执行。模型并行和数据并行同时支持。
  • Kunpeng[34]是阿里巴巴基于PS的分布式学习系统,支持像稀疏逻辑回归和多项式加性回归树这样的模型。

这些基于PS的框架通常设计用来支持传统模型,如逻辑回归或卷积神经网络,它们只包含稀疏部分或密集部分。这些框架很难处理涉及大规模稀疏特征和定制深度网络的模型,如[12, 31, 32]中所述。

另一方面,机器翻译[29]和语音识别[3]等自然语言处理系统使用类似的SparseNet + DenseNet模型结构。然而,它们通常只有多达百万个单词,而我们通常有数十亿个项目ID或图像作为特征。

另一个活跃的研究领域是自动寻找分布式模型训练的最佳运行时配置

  • [27]使用强化学习来寻找最优设备放置。
  • [19, 20]基于Legion[4]搜索混合并行策略的最佳配置,Legion是一种用于异构和并行机器的高性能编程系统。

然而,这些算法都没有考虑到稀疏性,也没有为如此规模的问题提供解决方案。

这些相关工作为构建高效的模型训练系统提供了许多见解。通过学习它们的优缺点,我们开发了XDL以支持涉及高维稀疏数据的工业问题。

3 XDL架构

在本节中,我们介绍XDL的架构和设计哲学。在讨论之前,有必要描述我们设计的动力。

3.1 涉及高维稀疏数据的深度模型的网络特性

近年来,基于深度学习的方法成功地革新了在线业务中应用的算法,并取得了最先进的性能。作为先驱工作,DSSM[17]提出在网络搜索引擎的场景中用深度网络建模查询和文档的相关性。DSSM模型的核心架构遵循embedding&多层感知器(MLP)范式:首先将高维sparse输入映射到低维embedding空间,然后使用多层感知器(MLP)拟合标签。这种embedding&MLP架构激发了大多数后续工作设计具有高维sparse输入的深度模型,这些模型来自视频推荐、网络搜索和电子商务网站广告等各种应用。代表性网络包括Wide & Deep Learning[8]、Youtube深度模型[10]、DeepFM[15]、深度兴趣网络(DIN)[32, 33]和CrossMedia[12]等。关于推荐系统深度模型的详细调查可以在[31]中找到。

从算法的角度来看,基于embedding&MLP的网络抽象了这些工业模型家族,这些模型通常将从稀疏数据中学习分为两个步骤:

  • i)表示学习(representation learning),从高维sparse输入中捕获信息并将它们embedding到低维空间,
  • ii)函数拟合(function fitting),模型denseembedding表示和监督label之间的关系。

为了简单起见,我们称第一步的网络为SparseNet,第二步的网络为DenseNet。图1说明了这种抽象。

图片名称

图1 embedding和多层感知器(MLP)网络架构的抽象,主要包含SparseNet和DenseNet。SparseNet将原始的sparse输入(维度为D,其中D可扩展至数十亿)映射到低维表示(维度为d,传统上扩展至数千)。DenseNet学习拟合数据的函数

在工业应用场景中,新的深度模型应该快速发展以提高业务利润。因此,算法科学家希望训练系统易于理解和使用:新模型应该在脚本代码中以独立视角设计,并并行运行,将复杂的分布式训练细节隐藏在后台。基于这些考虑,我们设计了XDL,这是一个工业深度学习框架,旨在构建一个高性能系统来训练涉及高维稀疏数据的深度模型。我们将展示我们如何应对高维稀疏数据带来的挑战,并提出一个遵循桥接策略的优雅设计。

3.2 XDL的设计哲学和桥接方法论

从上述内容中,我们可以看到涉及高维稀疏数据的模型中的SparseNet和DenseNet需要不同的系统功能。DenseNet由几个密集层组成,需要在本地机器上进行高计算密度。包括Tensorflow、MxNet和Pytorch在内的几个DL框架已经开源,这些框架可以很好地处理这类DenseNet。然而,SparseNet包含数千亿个特征,这些特征来自原始样本。因此,能够处理这种情况的成功DL框架必须具有出色的分布式特性和足够的稀疏数据计算能力。上述提到的现有DL框架并不针对涉及高维稀疏数据的网络设计良好。

为了更好地理解训练涉及高维稀疏数据的深度模型的挑战,让我们详细看看SparseNet和DenseNet的规模。首先展示我们数据集的规模。表1显示了我们在展示广告系统中一天使用的典型生产数据量。

表1

表2显示了我们日常任务中使用的一些模型的网络参数统计。显然,SparseNet是这些模型中贡献最大困难的部分。DenseNet遵循现有深度学习框架所持有的传统设置。此外,SparseNet需要处理输入数据的实际I/O问题,其体积为数TB,以及复杂和繁重的并行问题。这些挑战使得SparseNet的训练与DenseNet相比截然不同且至关重要。

模型 输入维度 网络参数 输出维度
MLP 10亿 18亿 1440
DIN 10亿 18亿 33438
CrossMedia 2.5亿 5亿 1464
模型 输入维度 网络参数 输出维度
MLP 1440 1.2万 1
DIN 33438 1.7万 1
CrossMedia 1464 1万 1

表2: 多层感知器、深度兴趣网络[32]和Crossmedia[12]模型的网络统计。

综上所述,我们设计了XDL,其架构如图2所示。遵循分而治之的策略,我们提出了一种全新的桥接架构,其中训练系统由两个主要子系统构成:

图片名称

图2 XDL架构。服务器和worker不需要具有相同的等级。它们可以部署在同一物理节点上,以提高通信效率。

  • 高级模型服务器(AMS)。AMS提供了一个精心设计的分布式系统,为训练SparseNet提供极高的效率和可扩展性。AMS处理快速演变的表示学习算法,这些算法具有大规模sparse输入。支持embedding字典和像CNN/RNN这样的模型,将大型sparse输入映射到密集向量。

  • 后端worker(BW)。BW遵循常见的深度学习设置,从低维输入中学习,并允许采用任何开源DL框架作为其后端。有了这种灵活性,我们的算法科学家可以轻松尝试涉及DNN[17]、注意力机制[32]或门控循环单元[33]的新模型结构,以模拟复杂且不断演变的用户行为。

在前向传递中,后端worker首先从I/O模块读取数据,然后向AMS发送特征ID请求。AMS将根据ID计算embedding向量并将其发送回后端worker。worker还将接收到上一次迭代的更新后的DenseNet模型参数。在收集完所有dense输入向量和模型参数后,worker将对输入向量执行sum pooling,并将结果输入到Dense中。在反向传递中,数据流被反转,worker计算梯度。AMS将收集DenseNet和SparseNet的梯度。参数更新将在服务器端使用像Adam[22]或Momentum SGD这样的求解器进行。关于AMS模块的实现和优化的更多细节将在下一节中描述。

XDL具有高性能和良好的可扩展性,因此能够支持工业生产。从2016年开始,XDL被用于训练目标广告部门的CTR预测模型。在数百台机器上运行,XDL训练了一系列具有数十亿参数的DNN模型,支持数百种业务场景

为了便于在各种计算平台上部署,XDL可以被多种资源管理平台调度,如Yarn,并为各种数据存储系统提供数据I/O接口,如HDFS和Kafka。XDL的良好兼容性使得业务应用能够快速升级。

3.3 高级模型服务器

从参数服务器的架构设计中学习,我们设计了高级模型服务器(AMS),这是一个管理SparseNet和DenseNet训练的分布式系统。AMS负责存储和更新所有参数,包括SparseNet和DenseNet。因此,参数放置(Parameter Placement)成为高性能设计的关键时刻。考虑到在线业务的稳定性和可用性,容错也应得到充分考虑。

参数放置(Parameter Placement)

由于sparse参数和dense参数都存储在AMS上,参数放置(Parameter Placement)算法应处理这两种参数的不同特性:

  • 一方面,sparse参数占用大量内存,因为稀疏特征的维度可能高达数千亿。然而,每个mini-batch中,只有少数特征ID被worker请求,这表明sparse参数的I/O压力很低。此外,随着训练过程的进行,sparse参数的内存使用量会发生变化,因为样本会不断将新的ID带入sparse参数中。
  • 另一方面,dense参数占用的内存较少,但I/O压力较大,因为worker在每个mini-batch中都请求整个dense参数。

根据上述分析,我们的kv存储使用hash mapping作为AMS中sparse参数的低级数据结构,这也存在于传统的参数服务器中。因此,我们可以在每个服务器上平均分配哈希表桶,减轻服务器上的内存存储压力。不幸的是,这个过程会导致worker请求数量大幅增加。为了解决这个问题,我们在embedding字典查找之前,在worker端合并worker请求。这种优化在实践中为我们提供了巨大的加速。至于dense参数,每个参数的I/O压力将在内存限制内平均分配在每个服务器上。

容错

随着样本数量和模型复杂性的快速增长,特定深度模型的离线训练需要越来越多的时间。当DL系统中的某些角色失败时,从一开始就重新训练在时间和资源方面都非常昂贵。此外,在一些在线学习场景中,稳定性和可用性非常重要。因此,XDL中精心设计了容错机制。AMS由调度器和多个服务器组成。服务器使用定期心跳与调度器保持同步。在训练过程中,整个模型的快照将存储在某个位置。如果任何服务器由于某种原因失败,调度器将注意到异常状态,并设置整个AMS系统为不可用状态,直到服务器被Yarn或其他调度系统重启。如果AMS准备好了,调度器将通知服务器从最新的快照恢复并继续训练过程。

同时,由于worker负责读取样本,每个worker的读取状态也作为特殊参数存储在AMS中。当worker失败并重新启动时,它们可以通过从AMS拉取参数来恢复其读取状态。通过这种努力,worker在训练过程中不会丢失或重复读取任何样本。

异步更新

XDL支持同步模式和异步模式。

  • 在同步模式下,worker在每次迭代中拉取参数和推送更新各一次。每次迭代直到所有AMS从所有worker接收更新并对参数应用平均梯度后才完成
  • 在异步模式下,没有这样的迭代概念。worker独立操作。因此,每个worker可以在将更新推送到AMS后立即进入下一步,而不需要等待其他worker完成他们的更新。同时,在服务器端,来自worker的更新以无锁风格应用。我们使用这种方法,因为embedding表的更新非常稀疏,冲突很少发生。对于密集DNN,尽管每个更新可能不会成功应用于整个模型,但总是只有一个全局模型版本。因此,无锁风格更新可以被视为在模型上添加一个掩码以随机丢弃部分梯度

在实践中,我们发现异步模式可以实现更大的系统吞吐量,同时几乎不会损失模型性能。

4 系统实现和优化

4.1 I/O

在生产场景中,我们训练系统的输入,如样本和初始模型,通常存储在不同的位置。例如,一些样本从在线队列系统中流式传输,而其他样本则从离线分布式文件系统中检索。I/O模块设计的目标是:最大化来自数据源(如HDFS、Kafka等)的输入带宽使用

众所周知,分布式ML系统中的I/O带宽已成为最大的瓶颈之一。由于我们使用现有的DL框架作为计算后端,如何更快地吞吐数据成为一个更关键的问题。因此,我们实现了各种优化以最大化I/O吞吐量。

首先,不同样本之间存在大量重复特征。这些重复来自我们数据源的最开始:用户行为。一个用户点击两个不同的链接将产生两个包含相同用户重复特征的不同样本。图3说明了重复的结构。压缩可以优化存储、通信和计算效率

图片名称

图3 重复

分层样本压缩:

XDL充分利用样本之间存在ID重复的事实。在预处理阶段,原始样本被组织成多前缀树(multi-prefix trees),这大大减少了存储空间和通信成本。在训练阶段,许多前缀树被组合成用户定义的batch大小的mini-batch。这种压缩也为计算带来好处,因为它减少了重复的embedding向量计算。如果前缀树的总样本大小超过batch size,最后一个前缀树将被分割。

图4展示了分层样本压缩的一个例子。原始样本被组织成两个不同的3层前缀树。树的第一层、第二层和第三层分别代表用户特征、广告特征和创意特征。在训练阶段,最后一个前缀树根据batch size=6进行分割。构建了两个辅助张量(指示器),以指示相邻层之间的关系。

图片名称

图4 I/O数据和计算压缩。在数据预处理阶段,生成多前缀树以用于训练阶段的mini-batch生成。

4.2 workflow流水线

workflow流水线是加速程序运行过程的常用技术。通过多线程技术使程序中的阶段在时间线上重叠,可以节省相当的运行时间。XDL中也实现了几种workflow流水线。

为了最大化I/O吞吐量,我们用线程流水线处理过程。一个完整的训练迭代被分为三个阶段:

  • 1)读取样本并将它们组合成mini-batch;
  • 2)从键值存储或输入数据中预提取mini-batch的参数索引,以在服务器上进行计算;
  • 3)拉取模型参数并进行前向/反向传播。

流水线利用线程并行性通过重叠执行这三个阶段。

如图5所示,这三个阶段的工作被调度到三个不同的线程池。这三个线程池通过一个无锁队列进行协调。如果使用GPU,除了I/O实现中的流水线外,还设计了SparseNet和DenseNet之间的workflow流水线,以进一步提高XDL的性能。在这种情况下,由于DenseNet的输入是SparseNet的输出,SparseNet的结果可以在DenseNet的前一个训练过程结束之前预计算。假设SparseNet和DenseNet的训练时间分别为$T_{sparse}$和$T_{dense}$,总运行时间将从$T_{sparse} + T_{dense}$减少到$max{T_spar se, Tdense}$。在这种情况下,训练过程的效率显著提高,模型性能损失很小,因为当SparseNet的输出被预提取时,可能存在一些未更新的稀疏ID。因此,当相邻mini-batch之间有很少共同稀疏ID时,这种流水线通常被适应。

图片名称

图5 训练迭代由三个阶段组成,使用流水线来最大化并行性。

XDL还允许算法科学家用几行代码构建连接训练过程中任何阶段的流水线。深度模型可以被划分为任意数量的阶段(stage)。算法科学家在考虑性能和准确性时有多种选择,以在训练过程中构建流水线。

4.3 高级模型服务器的优化

为了提高AMS的效率,在各个方面都采用了优化。

为了加速embedding过程中的K-V查询,XDL利用了GPU的能力。GPU有两个优势:

  • 首先,GPU的内存带宽更高(Nvidia P100为700GBps,Intel E5-2699 v4为100GBps)。
  • 其次,它在数千个核心上并行运行(Nvidia P100有3584个核心)。

因此,我们可以在GPU上实现更快的embedding字典查找。GPU的缺点是GRAM大小有限。为了解决这个问题,只有embedding字典的索引将被预提取到GRAM中。

由于AMS在大多数情况下扮演参数服务器的角色,因此AMS和后端worker之间存在大量的网络通信。当并行度很高时,网络通信量可能非常大。在这种情况下,网络通信消耗的运行时间成为总运行时间的关键部分。XDL采用Seastar[28]作为其通信库,并在Seastar框架的基础上进行了许多有效的优化。涉及零拷贝和CPU绑定等技术,以确保网络通信不会成为整个训练过程的瓶颈。

在深度学习中,batch size是一个非常重要的超参数。一些模型在使用mini-batch size时表现出良好的性能,而一些模型更喜欢大batch。这要求DL框架能够处理不同的样本batch size。正如我们之前提到的,CPU绑定(cpu-binding)是AMS中的一项重要技术。AMS中的每个线程都绑定到一个CPU核心,因此节省了不同核心之间的线程切换时间。

  • 当样本batch size较小时,这种CPU绑定机制使AMS上的计算性能高。
  • 然而,如果样本batch size很大,将有太多的ID需要由单个线程处理。

在这种情况下,AMS有一个自适应机制,其中创建了一个额外的线程池,用于处理大量的ID。有了这种优化,AMS在mini-batch和大batch样本上都表现出色。

4.4 使用XDL进行在线学习

在线学习方法最近在工业中被广泛使用,因为它能够实时捕捉客户行为的变化。一个有效的在线学习模型在许多业务场景中非常有价值。XDL为在线学习提供了一系列机制。

特征入口过滤器(Feature Entry Filter)

正如我们之前提到的,由于样本将新的ID带入sparse参数,sparse参数的内存使用量随着训练过程而变化。因此,模型存储在AMS上不断增加。因此,必须将ID的规模控制在一定水平。为了解决这个问题,XDL提供了几种特征入口过滤器机制,如基于概率的过滤器和计数布隆过滤器。在这种情况下,低频ID将在训练过程之前被丢弃,因为它们对模型的贡献较小。同时,在线模型的规模和内存存储得到很好的控制,这保证了在线系统的稳定性。

增量模型导出

随着持续的训练过程,当模型存储达到数百GB甚至更高时,全模型导出在时间和计算资源方面将变得极其昂贵。因此,XDL为在线训练和推理系统提供了增量模型导出。每个增量模型只需要少量的存储空间,并且非常容易部署在在线系统上。

特征过期

与特征入口过滤器具有相同目的,特征过期机制也控制模型的大小。在在线训练过程中,长时间未更新的特征将变得无用,可以被删除。XDL允许算法专家编写自己的特征过期函数,并定制特征过期计划。

用户界面

XDL为算法科学家提供了Python API,以便开发新的深度学习模型。这些API让用户专注于模型本身,而不必担心I/O效率、分布式策略、系统一致性和低级通信。在并行化模型上的编码和调试几乎与在串行模型上相同。通过隐藏并行化和优化的细节,用户可以非常容易地将新开发的模型部署到沙箱和生产环境中的大量GPU集群上。大多数情况下,科学家只需要修改DenseNet结构来测试新想法。用户可以通过运行用户定义的代码完全控制后端worker。修改数据流和SparseNet组件以尝试更激进的算法创新也非常方便。

5 XDL生态系统

除了XDL,阿里巴巴还开源了XDL算法解决方案和Blaze推理引擎,它们与XDL框架一起构成了一个生态系统。所有这些基于XDL的工具都可以在https://github.com/alibaba/x-deeplearning找到。

XDL算法解决方案包括几个在阿里巴巴广告业务中成功采用的有效模型,如深度兴趣网络(DIN)[32]、深度兴趣演化网络(DIEN)[33]、基于树的深度匹配(TDM)[35]等。

Blaze推理引擎是一个高性能的推理引擎,它利用内核融合和低精度推理等技术,为在线广告业务启用了大规模实时计算。

XDL生态系统为广告行业树立了示例和指南。越来越多的公司采用XDL作为他们的离线或在线深度学习工具。

6 评估

在本节中,我们提供了前一节描述的优化策略的详细评估。实验在通过25Gbps以太网互联的节点上进行。每个节点配备了2个Intel(R) Xeon(R) Platinum 8163 CPU @ 2.50GHz(96核)和512GB RAM。在编译时,我们使用GCC 4.8.5,带有-O2 -fopenmp -mavx2标志。

在我们的实践中,当选择后端时,MxNet或TensorFlow的性能几乎相同,因为它们在DenseNet计算方面具有相同的性能。在以下评估中,我们选择TensorFlow作为XDL的计算后端。

6.1 数据集

在我们的实验中,我们使用[25]发布的开放数据集,训练了一个用于点击率(CTR)预测问题的深度神经网络(DNN)模型。该数据集是从淘宝的电子商务推荐系统中以1%的比例抽样的。该数据集的统计信息显示在表3中。

开放数据集包括三个方面:用户兴趣、广告特征和上下文。它们中的每一个都由高维稀疏ID表示。使用XDL的Python接口,我们基于SparseNet + DenseNet抽象构建了模型

  • 在SparseNet中,我们使用embedding字典将每个ID转换为维度为18的密集向量。
  • 随后是一个分组sum pooling,将这些向量减少到23个向量。
  • 在DenseNet中,我们使用一个3层全连接DNN,包含约100K参数。(在这个DNN中,输入层有23*18个节点,第一层有200个节点,第二层有80个节点,输出层有1个节点。)

模型使用点击label来估计CTR。

6.2 子系统评估

为了了解XDL的性能,我们将逐一评估优化策略,看看我们能从每一个策略中获得多少好处。以下实验在上述开放数据集上使用相同的环境运行。所有实验在200个worker(除了可扩展性实验)上运行,使用80个CPU模式的AMS。我们为每个worker分配8个CPU核心,每个AMS将占用一台机器中的所有96个CPU核心。XDL在异步模式下运行,更新是无锁的,如第4.3节所述。

分层样本压缩:

I/O模块中的分层样本压缩对通信和计算都带来了巨大的加速。平均来说,每个用户ID在开放数据集中可以重复57次,这意味着那些用户共享的公共特征的压缩比可以达到57。为了评估使用样本压缩可以获得多少性能提升,我们重新抽样数据集以构建具有不同平均压缩比的新数据集。然后我们比较XDL在这些数据集上运行时的吞吐量。结果可以在图6中看到。

图片名称

图6 压缩比与吞吐量。运行在200个工作进程上,batch size为5000。

大batch size

使用大batch size将减少迭代次数,因此可以大幅降低数据复制和通信的开销。尽管Keskar等人[21]发现使用大batch训练可能导致尖锐的最小值并损害泛化。已有几种技术[13, 16]被提出来缓解这个问题。因此,我们倾向于使用更大的batch size来提高吞吐量。不同batch size下AMS的吞吐量显示在图7中。

图片名称

图7 批量大小与吞吐量。压缩比为57,使用200个工作进程。

可扩展性

我们还测试了XDL的可扩展性,并将其与不同数量的worker上的原生TensorFlow进行了比较,使用了80个AMS。同样,每个worker使用8个CPU核心,每个AMS使用96个CPU核心。XDL和TensorFlow都使用无锁风格的异步更新。结果显示在图8中。由于其对通信和重复特征的优化,XDL始终表现更好。当我们固定AMS的数量时,在异步模式下运行时,我们几乎实现了线性可扩展性。

图片名称

图8

7 结论

我们介绍了一个名为XDL的深度学习框架,它对于具有SparseNet + DenseNet抽象的高维稀疏数据集非常强大。凭借面向高维度的设计和系统优化,XDL在真实世界的数据集上可以比原生TensorFlow快得多。我们的算法科学家可以轻松地使用XDL灵活的API开发新模型,以支持阿里巴巴业务的快速发展。

随着XDL生态系统的开源进程,XDL可以在未来进一步发展。首先,像TVM[7]、深度梯度压缩[24]和混合精度训练[26]这样的高级加速技术正在测试中,并将很快部署。其次,作为参数存储和更新的开放解决方案,AMS将需要更灵活的数据流和模型结构,以满足最新模型如记忆网络和动态计算图的要求。

参考

介绍

baidu在2020《AIBox: CTR Prediction Model Training on a Single Node》提出了它们的AIBox平台。我们看下它的实现:

摘要

作为世界上主要的搜索引擎之一,百度的赞助搜索(Sponsored Search)自2013年起就采用了深度神经网络(DNN)模型来预测广告点击率(CTR)。百度在线广告系统(即“凤巢(Phoenix Nest)”)使用的输入特征极其高维(例如,成百上千的十亿个特征)且极其稀疏。百度生产系统中使用的CTR模型的大小可以轻松超过10TB。这给在生产环境中对这些模型进行训练、更新、使用带来了巨大的挑战。

对于百度的广告系统来说,保持模型训练过程的高效率显然是重要的,以便工程师(和研究人员)能够快速完善(refine)和测试他们的新模型或新功能。此外,由于每天有数十亿的用户广告点击历史记录条目到达,模型必须快速重新训练,因为CTR预测是一项极其时间敏感的任务。百度当前的CTR模型是在MPI(消息传递接口)集群上训练的,这需要高容错性和同步,从而产生昂贵的通信和计算成本。当然,集群的维护成本也是巨大的。

本文介绍了AIBox,这是一个集中式系统,通过使用固态硬盘(SSD)和GPU来训练具有数十TB规模参数的CTR模型。由于GPU上的内存限制,我们仔细地将CTR模型划分为两部分:一部分适合CPU,另一部分适合GPU。我们进一步引入了一个双层缓存管理系统,用于在SSD上存储10TB参数的同时提供低延迟访问。在生产数据上的广泛实验显示了新系统的有效性。AIBox具有与大型MPI集群相当的训练性能,同时只需要一小部分集群的成本。

1 引言

作为世界上领先的搜索引擎之一,百度的赞助搜索系统(即“凤巢”)[20] 自2013年起就采用了深度神经网络(DNN)模型来预测广告点击率(CTR)。CTR预测[19, 22]在决定最佳广告空间分配中起着关键作用,因为它直接影响用户体验和广告盈利能力。通常,CTR预测以多个资源作为输入,例如查询-广告相关性、广告特征和用户画像。然后,它估计用户点击给定广告的概率。最近,深度学习在计算机视觉[29]和自然语言处理[5]中取得了巨大成功。受此启发,提出了用于CTR预测任务的学习方法[13, 15, 20, 61, 65, 66]。与常用的逻辑回归[40]相比,深度学习模型可以显著提高准确性,但训练成本也显著增加。

在百度搜索广告的当前生产系统中,我们模型的训练过程既耗费资源又耗时。模型是在MPI(消息传递接口)[52]集群中的参数服务器[24, 32]上训练的,该集群拥有数百个CPU-only节点。生产中使用的主要模型大小超过10TB,存储/管理在特殊硬件上。参数服务器解决方案受到百级节点环境中节点故障和网络故障的影响。更糟糕的是,参数服务器中的同步阻塞了训练计算,并产生了大量的网络通信开销,而异步训练框架由于每个工作节点上的模型过时[48, 62],存在模型收敛问题

改进赞助搜索的生产系统存在许多不同的方向上的迷人机会和挑战。一个活跃的研究领域是在调用CTR模型之前提高“召回”(广告)的质量。例如,百度最近向社区分享了这样的技术论文[20],该论文建立在快速近邻搜索算法[55, 64]和最大内积搜索技术[51, 56]之上。

在本文中,我们介绍了百度为改进在线广告系统的另一项并行主要努力,即将CTR模型训练从MPI集群转移到GPU上。虽然使用GPU进行机器学习和科学计算已经成为一种常见做法,但使用GPU训练商业CTR模型仍然带来许多实质性挑战。最值得注意的挑战是训练数据规模达到PB(PeteByte)级别,训练出的模型大小超过10TB。训练样本的数量可以达到数千亿,特征的数量可以达到数千亿(我们通常使用2的64次方作为一个方便的特征空间大小的替代)。输入模型的数据也非常稀疏,每个特征向量仅有几百个非零项。

作为商业赞助搜索系统,任何模型压缩技术都不应妥协预测性能(收入)。实际上,即使是微小的(例如,0.1%)预测准确性下降也会导致不可接受的收入损失。事实上,整个系统已经高度优化,几乎没有冗余(例如,参数已经被仔细量化为整数)。看来改进的空间非常小。流行的模型压缩技术,如下采样[8, 30, 46]和哈希[12, 27, 33, 34, 57],对于训练数据极其高维(例如,数百亿个特征)和极其稀疏(例如,每个特征仅有几百个非零项)的商业CTR模型来说效果较差。研究论文中常见的论点,如“以仅损失0.3%的准确性为代价减少一半的训练成本”,在这个行业不再适用。另一方面,在百度训练DNN CTR模型是一项日常任务。工程师和数据科学家必须尝试许多不同的模型/特征/策略/参数,并且必须非常频繁地训练和重新训练CTR模型。硬件(如MPI集群)和能源消耗的成本可能非常高。

为了应对这些挑战,我们提出了AIBox,这是一个新颖的集中式系统,可以在单个节点上高效地训练这个庞大的机器学习模型。AIBox采用新兴硬件,SSD(固态硬盘)和GPU,存储大量参数并加速繁重的神经网络训练计算。作为一个集中式系统,AIBox直接消除了分布式系统中由网络通信引起的缺点。此外,与大型计算集群中的数千个节点相比,单个节点的AIBox硬件故障数量要少几个数量级[49]。由于只需要内存中的锁和GPU片上通信,单个节点中的同步成本也大大降低。与分布式环境相比,没有通过网络传输数据。

然而,在设计AIBox时存在两个主要挑战。

第一个挑战是:在单个节点上存储10TB规模的模型参数。当容量超过1TB时,内存价格飙升。随着模型未来变得更大,这在可扩展性上是不切实际的,并且在实际应用中大规模生产也不实用。由于成本高昂,我们不能将整个10TB参数存储在主内存中。新兴的PCIe总线上的非易失性内存高速(NVMe)SSD比硬盘驱动器的延迟低50倍以上[60]。我们利用SSD作为二级存储来保存参数。然而,SSD有两个缺点:

  • 首先,在延迟方面,SSD仍然比主内存慢两个数量级,导致训练过程中参数的访问和更新速度慢。
  • SSD的另一个缺点是SSD中的存储单元仅能持续几千次写入周期。

因此,我们必须维护一个有效的内存缓存,以隐藏SSD延迟并减少对SSD的磁盘写入。

第二个挑战是:在单个节点上使用多个GPU来加速训练计算。最近,配备32GB高带宽内存(HBM)的Nvidia Tesla V100[28]的单精度性能达到15.7 TFLOPS,比顶级服务器CPU节点(Intel Xeon系列)在深度学习推理上快47倍[38]。这为设计一个多GPU计算节点提供了独特的机会,该节点的计算性能与集群相当。然而,当前现成的GPU没有TB级HBM。我们不能将整个CTR预测神经网络保持在GPU HBM中。在这项工作中,我们提出了一种新颖的方法(第2节)将神经网络划分为两部分:

  • 第一部分是内存密集型的,并且在CPU上训练。
  • 网络的另一部分是计算密集型的,同时输入特征数量有限,我们在GPU上训练它。

训练数据和模型参数在主内存和多GPU的HBM之间传输。然而,主内存和GPU HBM之间的通信受到PCIe总线带宽的限制。当通信带宽成为瓶颈时,高GPU数值计算性能被阻塞。新兴的NVLink[21]和NVSwitch[44]技术使GPU之间的直接通信无需涉及PCIe总线。我们采用NVLink并设计了一个HBM内的参数服务器,以减少GPU数据传输。

总结来说,本工作的主要贡献包括:

  • 我们介绍了AIBox,这是一个由SSD和GPU加速的单节点系统,用于训练具有10TB参数的CTR预测模型。单节点设计范式消除了分布式系统的昂贵网络通信和同步成本。据我们所知,AIBox是第一个为如此大规模的实际机器学习应用设计的集中式系统。
  • 我们展示了一种新颖的方法,将大型CTR预测模型划分为两部分。划分后,我们能够在CPU上保留内存密集型训练部分,并利用内存受限的GPU来加速计算密集部分的训练
  • 我们提出了稀疏表(Sparse Table)来通过在SSD上存储模型参数并利用内存作为快速缓存来减少SSD I/O延迟。此外,我们实现了一个3阶段流水线,将网络、稀疏表和CPU-GPU学习阶段的执行重叠起来。
  • 我们通过将AIBox与由75个节点组成的分布式集群解决方案进行比较,在包含10PB样本的实际CTR预测数据上进行了广泛的实验评估。结果表明AIBox的有效性——AIBox具有与集群解决方案相当的训练性能,而其成本仅为我们为集群支付成本的不到1/10。

本文的其余部分组织如下。第2节介绍了我们的CTR预测神经网络模型。第3节,我们介绍了AIBox的高层设计。我们在第4节探讨了在稀疏表上大型模型存储的挑战,并提出了采用SSD上的双层缓存管理系统的解决方案。第5节展示了实验结果。第6节讨论了相关工作,第7节总结了本文。

2 CTR预测神经网络

工业深度网络被设计和训练以处理大规模数据样本,以帮助准确、快速、可靠地预测广告的点击率(CTR)。百度CTR预测模型中的特征通常是极其稀疏的特征(例如,成百上千的十亿个特征),每个向量中仅有极少数非零值(例如,几百个)。这个庞大的DNN模型在仅存储非零参数并进行仔细量化后,参数大小超过10TB。由于GPU的HBM容量有限,显然不切实际将整个模型的10TB参数保持在GPU的HBM中

在本文中,我们提出了一个双模块架构,用于在CPU+GPU上训练庞大的DNN CTR模型。第一个模块专注于高维&稀疏特征的嵌入学习,第二个模块用于与第一个模块产生的密集特征进行联合学习。嵌入学习在CPU上处理,以帮助学习低维密集嵌入表示。由于10TB参数的内存密集问题使得在训练期间无法将整个模型保持在内存中,我们利用SSD存储模型参数。参数可以从SSD快速访问到CPU。

通过将从CPU学习到的嵌入向量传输到GPU,计算密集型的联合学习模块可以充分利用强大的GPU进行CTR预测。在联合学习模块中,几个全连接神经网络被建模,以嵌入作为输入。这些神经网络的最后一层被连接在一起,用于最终的CTR预测。图1显示了设计的CTR神经网络模型的概览。我们在以下小节中介绍模型的细节。

图片名称

图1: 设计的CTR预测神经网络模型概览。嵌入学习输入层的节点:代表高维稀疏特征。联合学习合并层中没有入箭头的节点是:除了嵌入之外的密集个性化输入特征。

为了方便讨论,在本文的其余部分,我们将使用$10^{12}$(千亿)作为CTR模型中极其稀疏特征的维度。读者应记住,实际生产系统中的特征数量可能远远超过$10^{12}$。

2.1 在CPU上的嵌入学习

嵌入学习模块旨在将高维稀疏向量(例如,$10^{12}$维)映射到低维密集表示。如图1所示,嵌入学习模块包括高维稀疏特征的输入层和输出嵌入层。使用ReLU作为激活函数。这个模块主要是内存密集型的,因为$10^{12}$特征导致10TB规模的模型参数,并且无法将所有参数加载到主内存中。为了学习嵌入,我们将10TB参数存储到SSD中。由于SSD和CPU之间的高效访问速度,我们可以轻松地从SSD加载参数并在CPU上学习embedding。

2.2 在GPU上的联合学习

在CPU上计算高维稀疏特征的嵌入后,我们将嵌入从CPU传输到GPU进行CTR预测过程。联合学习的输入包括:dense个性化特征和学习到的embedding。个性化特征通常来自多种来源,包括广告创意的文本、用户个性化行为和各种与广告相关的元数据。如果我们直接将这些特征进行concate并输入到神经网络中,可能无法充分探索个性化特征中的信息,导致CTR预测结果不准确。因此,我们设计了几个深度神经网络,并联合学习有意义的表示,以进行最终的CTR预测。如图1所示,联合学习模块包含几个(图中为两个)深度神经网络。每个网络将学习到的嵌入和一种个性化信息一起作为输入层。然后应用几个全连接层,以自动方式帮助捕获特征的交互。这些网络的最后一层隐藏层被组合用于softmax层和CTR预测的输出层。我们使用负对数似然作为目标函数:

\[L = -\frac{1}{|S|} \sum_{(x,y) \in S} (y \log p(x) + (1 - y) \log (1 - p(x)))\]

这里:

  • $ S $是训练集,
  • $ \mid S \mid $是训练样本的大小。
  • $ x $代表我们CTR预测模型的输入特征,
  • $ y \in \lbrace 0, 1 \rbrace $是标签,表示用户是否点击了广告。
  • $ p(x) $是softmax层之后的输出,表示数据样本被点击的预测概率。

为了有效地从前一个神经网络中学习,表示从第一层和最后一层隐藏层提取出来,然后与当前神经网络的输入层连接,用于联合学习。具体来说,第一隐藏层代表低级特征学习,从输入层提取最相关的信息。最后一层隐藏层显示高级特征学习,并检测对最终CTR预测最抽象但最有帮助的信息。我们结合从前一个网络中最有意义的低级和最强大的高级信息,以获得更准确的CTR预测结果。

这个模块主要是计算密集型的,因为有多个神经网络。我们在GPU上训练这个模块以加速。完成一轮或几轮联合学习后,使用反向传播首先更新GPU上的神经网络参数和传输的嵌入。然后更新的嵌入被返回以更新CPU上的模型参数。通过这种方式,两个模块共同工作进行CTR模型训练。

3 AIBox系统概览

在本节中,我们介绍AIBox系统的概览,并从高层次描述其主要模块。图2描绘了AIBox的架构。它包含三个组件:CPU模块、稀疏表模块和GPU模块。

图片名称

图2 AIBox架构

CPU模块:协调和嵌入学习

CPU模块协调整个训练工作流程。首先,它通过高速网络从分布式文件系统(例如,HDFS [7])读取训练样本,并构建后续并行处理的mini-batches。我们称这些mini-batches为一个pass。然后,CPU模块通过与稀疏表模块交互来计算mini-batches的特征嵌入,以获取引用的参数。之后,嵌入被传输到GPU进行神经网络的联合学习部分。在GPU上完成联合学习后,GPU模块内联合神经网络的反向传播梯度返回到CPU模块。它通过稀疏表的push操作更新CPU侧的参数。除了主工作流程外,CPU模块还会定期将当前训练快照保存为检查点到分布式文件系统,以便作业失败恢复。

稀疏表:关键值存储系统

稀疏表是一个键值存储系统,存储SSD上的$10^{12}$离散特征的值(例如,模型权重和点击信息)。这些特征消耗超过10TB的空间。特征到文件映射的键哈希索引位于内存中进行键查找。尽管SSD的延迟比硬盘低得多,但它们仍然无法与纳秒级内存延迟相媲美。为了解决这一关键问题,我们为SSD访问设计了一个显式的内存缓存策略以隐藏延迟。同时,缓存还充当缓冲区,避免了频繁写入SSD(第4节)。

GPU模块:联合学习

GPU模块接收从CPU模块传输的sparse embedding,并将它们存储在HBMs中。然后,embedding被送入dense联合学习网络。通过PCI-E总线传输的embedding和GPU核心上的联合学习计算是时序重叠的。我们为数据传输创建一个单独的CUDA流[43],另一个CUDA流用于学习计算。GPU HBMs充当片上参数服务器(on-chip parameter server)。dense神经网络的参数在所有GPU之间复制,并存储在片上GPU HBM中。对于每个mini-batches的pass,每个GPU根据其本地副本计算新参数。之后,AIBox中的所有GPU通过NVLink高速互连进行集体通信以同步参数。

3阶段流水线

训练工作流程主要涉及3个耗时任务:训练样本读取、sparse表操作和神经网络训练(嵌入学习+联合学习)。这3个任务对应于独立的硬件资源:网络、SSD和CPU+GPU。如图3所示,我们构建了一个3阶段流水线来隐藏这些任务的延迟,并为每个阶段维护一个预取队列。此外,预取还隐藏了阶段之间的通信时间。对于每个阶段,我们创建一个工作线程,从预取队列中取作业并提供给相应的资源。然后,处理结果被推入下一阶段的预取队列。当下一阶段的预取队列已满时,即下一阶段已经有太多未处理的作业时,每个阶段的线程将被暂停。预取队列的容量根据每个阶段的执行时间预设。对于图3中的RE阶段,我们没有任何执行依赖。因此,我们通过网络从分布式文件系统中提取样本,直到我们在队列中有足够的未处理样本为止。ST阶段仅依赖于RE——通过我们的缓存管理系统消除了对EL和JL的依赖。ST阶段从SSD中加载RE中引用的参数。由于我们的缓存内存足够大,可以缓存过去几次pass中引用的参数,我们不需要等待EL和JL的完成。我们只需要确保加载的参数在被EL + JL阶段的相应pass使用之前不能从缓存中移除。EL计算依赖于ST阶段,因为我们只能在将参数加载到内存后开始计算嵌入。而JL训练必须等待从EL部分计算出的嵌入。由于数据依赖性,JL和EL处于一个流水线中。配备多个高端GPU的AIBox具有与数十个仅CPU节点的分布式集群相当的计算能力。我们可以为神经网络模型的训练实现类似的执行时间。此外,RE是最简单的和最快的阶段。其延迟在流水线中被隐藏。然而,与所有参数都在内存中的分布式集群解决方案相比,与SSD交互的稀疏表操作纯粹是开销。我们需要积极优化ST,以确保ST比EL + JL更快。在这种情况下,ST的延迟在流水线中被隐藏,以便我们可以拥有相当的训练速度。因此,我们关注以下文章中稀疏表的设计细节。

图片名称

图3 设计的CTR预测神经网络模型包括一个3阶段流水线,该流水线对于网络、SSD和神经网络训练(CPU + GPU)上的执行操作是时序重叠的。

4 稀疏表

稀疏表旨在高效地在SSD上存储模型参数。它利用内存作为SSD的快速缓存,同时减少SSD I/O并提供低延迟的SSD访问。它由两个主要组件组成:键哈希索引和双层缓存管理。

4.1 键哈希索引

为了通过特征键访问SSD上的参数文件,我们必须为CTR预测模型中的$10^{12}$个参数存储$10^{12}$个键到文件的映射。在内存中将每个键到文件的映射存储为64位值对需要1.6TB = (8字节键 + 8字节SSD上的偏移量) × $10^{12}$,这超出了1TB的内存预算。我们必须仔细设计键哈希索引和SSD上的文件结构以减少内存占用。

我们引入了一个分组函数,将键映射到组ID,使得每个组包含$m$个键,即group(key) → {0, 1, ···, $10^{12}/m − 1$}。这里$10^{12}$个键被划分为$10^{12}/m$组。在对键进行分组后,我们能够保持组到文件的映射在内存中,因为内存消耗仅为原始键到文件映射的$1/m$。由于键从1到$10^{12}$是连续的,分组函数可以通过均匀划分键空间轻松获得,例如,group(key) → key mod $10^{12}/m$。我们设置$m = \lfloor BLOCK/(8 + \text{sizeof(value))} \rfloor$,其中BLOCK是SSD的I/O单元,它由SSD块大小(通常是4096)决定,8代表键占用的字节,sizeof(value)是值(模型参数)的大小,字节为单位,在我们的CTR预测模型中大约是50字节。$m$绝不能设置为小于$\lfloor BLOCK/(8+\text{sizeof(value))} \rfloor$的值,因为SSD访问必须从磁盘获取BLOCK字节。设置一个太小的$m$是次优的。另一方面,我们选择的$m$越大,键哈希索引的内存占用就越小。然而,一个较大的$m$导致一个较大的组,因为我们必须从SSD获取多个页面以获得一个组。因此,我们设置的$m$值在组到文件映射占用内存的可接受空间时是最优的。当块大小远大于值大小时,这是真的。

作为内存占用的权衡,分组策略的缺点是同一组中的值从磁盘获取,即使它们在当前mini-batch中未被引用——I/O带宽被浪费。一个可能的优化是将具有高共现特征的组在一起,例如,预训练一个学习哈希函数[26]以最大化特征共现。这属于垂直分区[42, 63]的另一个研究领域,超出了本文的范围。此外,这个缺点被缓存管理组件减少,我们跳过缓存键的组从磁盘读取。

4.2 双层缓存管理

缓存管理设计受到以下两个挑战的指导:访问性能和SSD的寿命。

首先,内存访问延迟是纳秒级,而SSD需要微秒级才能查看数据,因为SSD比内存慢大约1,000倍。然而,CTR中的参数如此稀疏和偏斜,以至于在mini-batch的一个pass中少于1%的参数被引用。这为我们提供了一个机会来构建一个内存中的缓存系统,保存有限内存预算中的频繁使用的“热参数”

第二个挑战是:SSD的物理属性只允许每个存储单元进行数千次写入。参数在训练的每次迭代中更新。如果参数立即更新,将显著缩短SSD的寿命。缓存管理还充当参数缓冲区。缓冲的参数在内存中更新,不涉及SSD I/O。当缓冲区达到其容量并且缓存替换策略将其交换出缓存时,它们才被懒洋洋地物化到SSD

我们提出了一个双层缓存管理来应对这些挑战(图4)。与通过链表链接处理冲突的经典哈希表相比,我们的系统在链接之前有一个额外的哈希级别。此外,我们引入了两个单独的链表来优化缓存探测性能。我们还为每个SSD文件附加了一个布隆过滤器,以减少不必要的SSD读取。一个后台工作线程运行以维护缓存策略,即执行缓存替换并将更新的参数写回SSD。

图片名称

图4 二级cache管理架构

第一层,图4中的Level 1,将键组ID哈希到缓存插槽$s_i—s_i = hash_1(g_{id})$,其中:$g_{id}$是在键分组中计算的组ID(g_id = group(key))。每个si都有一个指向SSD中的文件的指针,该文件存储所有处理过的键的序列化模型参数,哈希1(g_id) = si。由于我们模型的极端稀疏性,许多参数在训练的早期阶段没有被触及。用默认值初始化所有参数并将其保存到SSD上是低效的。对于第一次引用的键,我们必须从SSD读取一个数据块,它只返回默认值。有$10^{12}$个参数在我们的模型中,它可能产生$10^{12}$个不必要的SSD读取。因此,我们执行延迟参数初始化,即我们直到它被mini-batch引用才物化一个参数。此外,我们为每个物化文件附加了一个布隆过滤器[9, 41]以过滤不必要的磁盘读取。布隆过滤器是一个低内存占用的概率数据结构,用于测试一个元素是否存在于一个集合中。它保证没有假阴性,因为当布隆过滤器返回假时,被测试的元素肯定不是集合的成员。此外,查询布隆过滤器是快速且内存带宽效率高的。只访问一个常数位数。

由于我们在布隆过滤器返回假时跳过探测SSD文件,布隆过滤器可能减少了SSD文件中潜在键缺失的数量。

第二层中的哈希码(图4中的Level 2)由hash2(g_id, bucket)计算,其中bucket是一个可调参数,控制每个缓存插槽si中的桶数。我们引入第二层哈希以缩短链表的链长。探测链表需要遍历链表直到我们找到探测键。在最坏情况下,它可能需要链表长度的步骤,而其他哈希表操作只需要恒定时间。因此,遍历链表是哈希表探测中最耗时的步骤。第二层哈希将si的长链链表分成多个桶,每个桶都有短链表以减少低效的迭代。参数bucket用空间换取高效探测。桶越大,我们期望中的链表就越短。

对于所有哈希到第二层同一个桶的键值对,我们构建两个单独的链表,最近最少使用[45](LRU)列表和最少频繁使用[53](LFU)列表,以在内存中缓存这些对(图4中的分离链表)。LRU列表存储当前pass的mini-batch中使用的参数。它将最近使用的键值对保持在链表的前面,从而减少遍历链表以定位数据的步骤。另一方面,LFU列表充当缓冲区,包含缓存替换的候选。最频繁使用的值位于链表的头部,因此可以高效访问。而最少频繁使用的值被冲入SSD并从缓存内存中移除。一个后台运行的工作线程定期将不在当前pass的mini-batch中的数据从LRU列表移动到LFU列表。链表导致许多小链表节点的分配和释放。我们维护一个内存池,采用Slab内存分配机制[6]来管理链表中节点的创建和删除。

4.3 稀疏表操作符

稀疏表的拉取(pull)操作符接收当前pass中引用的一组键作为输入,并将这些键的值加载到内存中。图5说明了如下工作流程:首先,我们通过键分组哈希函数计算组ID,并在双层缓存管理的第一层定位缓存插槽si。然后,我们查询每个定位缓存插槽的布隆过滤器,以测试这些引用的键可能存在(当键第一次被引用时不存在)。对于每个未通过布隆过滤器的键,我们在内存中创建一个带有默认值的链表节点,并将键插入布隆过滤器。对于通过布隆过滤器的键(它们可能存在于缓存或SSD文件中),我们通过它们的键和组ID在第二层定位相应的缓存桶bi,j,并遍历两个链表(LRU和LFU)以查找引用的键。如果在链表中找到键,则为缓存命中,因为我们不需要访问SSD。另一方面,对于未在链表中定位的键,我们必须读取SSD上的文件。以布隆过滤器理论上保证的低概率,引用的键在文件中未找到。我们为每个未找到的键创建一个默认节点,并像对未通过布隆过滤器的键一样更新相应的布隆过滤器。最后,我们将定位/创建的链表节点移动到LRU列表的头部。拉取操作后,所有引用键的参数都存储在内存缓存中,并且它们已准备好使用。

图片名称

图5

稀疏表的推送(push)操作符通过给定的一组键更新模型参数。更新通过学习率和反向传播梯度的组合值来增加引用参数。我们的工作流程(第3节)确保所有给定键的参数已通过相同键的拉取操作符加载到内存中。因此,我们可以通过双层缓存管理定位这些参数的值,并在不涉及写入SSD数据的情况下高效地增加它们。磁盘写入被推迟,直到更新的缓存节点从缓存内存中移除。当缓存内存预算已满时,即LFU列表的长度大于预设阈值MAX_LFU,我们从LFU列表中移除最后的FLUSH_STEP个节点,其中FLUSH_STEP是一个可调常数。由于节点按频率降序排序在LFU列表中,最后的FLUSH_STEP个节点的使用频率最低。被移除的参数必须被刷新到SSD,因为它们已被过去的推送操作更新。我们通过缓存管理组件的第一层定位包含被移除键的文件,从SSD读取它们,并将它们与更新的值合并到新文件中。

4.4 文件管理

推送稀疏表操作符不断在SSD上创建小文件,而现代操作系统上的本地文件系统,例如ext4[39]和XFS[54],并不是为处理许多小文件而设计的。大量的文件元数据压倒了这些文件系统。我们通过利用小文件创建总是成批进行的事实,为这项任务设计了一个特定的轻量级文件管理系统。与键分组类似,我们将F个小文件分组到一个大文件中以创建更少的文件。此外,当将一批独立的小文件磁盘写入分组为顺序写入时,它充分利用了磁盘带宽。对于每个小文件,其在大文件中的偏移量被保存。偏移量和大文件的名称的组合作为小文件的“文件名”。这种组合在缓存管理系统的第一层缓存插槽中维护。当我们将更新刷新到SSD时,它会更新。在我们创建越来越多的新文件后,一些旧文件没有被任何缓存插槽引用,就会过时。我们在内存中为每个大文件维护一个引用计数器。当一个“小文件”在其中更新并引用一个新文件时,计数器减少。我们一旦其计数器达到零就删除大文件。此外,我们定期监控磁盘使用情况。我们合并引用计数低的文件,并在磁盘使用量高于模型大小MAX_REPLICATION倍时删除它们。这里MAX_REPLICATION由SSD容量 ∗ (85%+超额配置)/模型大小计算,其中超额配置是制造商指定的SSD的保留存储空间。为了最大化SSD寿命,SSD控制器采用磨损均衡算法[11],这需要空闲空间来分散磁盘写入。当空闲空间少于85%时,SSD的性能会下降[2]。

5.实验

参考

Pinterest在《Graph Convolutional Neural Networks for Web-Scale Recommender Systems》中提出了PinSage:

介绍

深度学习方法已经在推荐系统应用中非常重要,被用于学习关于图片、文本、单个用户的有用高维embeddings。使用deep models学到的representations可以被用于补全、或者替换传统的推荐算法(比如:CF)。并且这些学到的representations具有高可用性,因为他们可以被用在许多推荐任务中。例如,使用一个deep model学到的item embeddings可以被用于item-item推荐,并且可以推荐有主题的集合(例如:playlists,或者“feed”内容)

最近,在该领域内有较大发展——特别是能够基于在图结构化数据上学习的新deep learning方法的发展,这对于推荐应用来说很基础(例如:可以利用user-to-item交互图,也可以使用社交网络图)。

在这些成就中,最著名的是Graph Convolutional Networks(GCNs)深度学习结构的成功。GCNs背后的核心思想是:学习如何从局部图邻居(local graph neighborhoods)的feature信息中进行迭代式聚合(aggregate)(图1)。这样的一个“卷积(convolution)”操作会从一个node的一跳图邻居的feature信息进行转换和聚合,并且通过将多个这样的convolutions信息进行stacking可以传播到图的远方。不同于纯content-based deep models(例如:RNNs),GCNs会同时利用内容信息和图结构。GCN-based方法已经在无数推荐系统benchmarks上设置了一个新的标准。然而,在benchmark任务上的收益还没有转换到真实生产环境中。

同时将GCN-based node embeddings的training和inference扩展到具有数十亿nodes和数百亿edges的图上是一个巨大的挑战。当在一个大数据环境中,对GCNs进行扩展很难,因为在设计底下的许多核心假设是冲突的。例如,在训练期间,所有已存在的GCN-based推荐系统需要在整个图拉普拉斯算子(full graph Laplacian)上操作,因而是不可行:当底层图(underlying graph)具有数十亿nodes,它们的结构通常是演化的。

当前工作

这里我们提出了一个高度可扩展的GCN框架,我们在Pinterest的生产环境中开发和部署过。我们的框架,是一个random-walk-based GCN 称为“PinSage”,它在30亿nodes和180亿edges的一个大型图(massive graph)上操作——该graph是GCNs的常见应用的10000倍以上大。

PinSage会利用许多关键insights来弹性提升GCNs的可扩展性:

  • on-the-fly convolutions:传统的GCN算法通过将feature matrics乘以full graph Laplacian的幂来执行图卷积。相反的,PinSage算法执行很高效,通过对在一个node周围的邻居进行抽样,并且从该抽样后的邻居来动态构建一个计算图。这些动态构建的计算图(图1)指定了如何在一个特定node周围执行一个局部卷积(localized convolution),并且缓和了在训练期间对整个graph操作的需求。
  • Producer-consumer minibatch construction:我们开发了一个producer-consumer架构来构建minibatches,它能确保在模型训练期间最大化GPU利用率。一个大内存、CPU-bound producer可以有效抽样node network邻居,并且获取必要。
  • 有效的MapReduce inference:给定一个fully-trained GCN模型,我们会设计一个有效的MapReduce pipeline,可以将训练好的模型分散来生成数十亿节点的embeddings,可以最小化重复计算。

图片名称

图1

除了在可扩展性上的这些基础优点外,我们也会引入新的训练技术和算法创新点。这些创新点会改善由PinSage学到的representations的质量,从而能导致在下游推荐系统任务上获得大的提升:

  • 通过random walks来构建convolutions:通过采用节点的所有邻居来执行convolutions(图1)会导致很大的计算图,因此,我们会采用sampling。然而,random sampling是次优的,我们会开发一种新的技术,它使用short random walks来采样计算图(computation graph)。一个额外的好处是:每个node具有一个importance score,我们会在pooling/aggregation step中使用。
  • Importance pooling:graph convolutions的一个核心组件是:在graph中的局部邻居的feature信息的aggregation。我们会引入一个方法来对在该aggregation中的node features的importance进行权衡,它基于random-walk相似度measures,在离线评估指标中会有一个46%的效果增益。
  • 课程培训(Curriculum training):我们设计了一个Curriculum training的scheme,其中该算法会在训练期间feed越来越hard的样本,从而产生一个12%的效果增益。

对于在Pinterest中的多个推荐任务,我们已经部署了PinSage,它是一个用户可以对交互的pins进行流行内容发现和管理的应用,其中:pins是在线内容的一个可视化书签。用户会将这些pins组织成boards,它包含了许多相似pins的collections。总之,Pinterest是世界上关于图片的最大user-curated graph(用户组织图),具有20亿唯一的pins,被收集到超过10亿的boards上。

通过大量离线metrics、A/B tests,我们展示了:对比起其它可扩展的deep content-based推荐算法,我们的方法在一个item-item推荐任务上(例如:相关pins推荐)以及一个”homefeed”推荐任务上可以达到SOTA的效果。在离线ranking指标上,我们对比baseline获得了超过40%的提升,在head-to-head人工评估上,我们的推荐要好60%。。

据我们所知,这是最大的deep graph embeddings,为新一代基于graph convolutional结构的推荐系统铺平了道路。

2.相关工作

3.方法

在本节中,我们会描述PinSage结构、training、以及MapReduce pipeline的技术细节,可以使用一个训练好的PinSage模型来有效生成embeddings。

我们的方法的关键计算重任是:局部图卷积(localized graph convolutions)。为了为一个node(例如:一个item)生成embedding,我们会应用多个convolutional模块,它们会聚合来自该node的局部图邻居(local graph neighborhood)(图1)的feature信息(例如:可视化features、文本features)。每个module会学习如何从一个小的图邻居(graph neighborhood)来聚合信息,通过将多个这样的模型进行stacking,我们的方法可以获得关于局部网络拓朴的信息。重要的是,这些localizied convolutional modules的参数会跨所有nodes进行共享,使得我们的方法的参数复杂度完全取决于input graph size。

3.1 问题设定

Pinterest是一个内容发现应用,其中:

  • 用户会将这些pins组织成boards,
  • 用户会将比较相关的pins包成集合

总之,Pinterest graph包含了20亿pins,10亿boards,并且超过180亿的edges(例如:pins的成员和它们相应的boards)。

我们的任务是,生成pins的high-quality embeddings 或者 representations(例如:通过最近邻查询相关的pin推荐,或者在下游的reranking系统中使用)。为了学习这些embeddings,我们会将Pinterest环境建模成一个二部图,它包含了两个不相互交集合I(包含了pins)和C(包含了boards)的结点。注意,我们的方法也是天然泛化的,其中I被看成是一个items集合,C被看成是user-defined contexts和collections。

除了图结构外,我们也会假设:pins/items \(u \in I\)会与real-valued属性相关联,\(x_u \in R^d\)。总之,这些属性可以指定关于一个item的metadata或content信息,另外在Pinterest的case中,我们有:具有丰富文本信息和图片features有关的pins。我们的目标是,利用这些input属性,以及二部图结构来生成高质量embeddings。这些embeddings接着会被用于通过最近邻lookup的推荐系统的候选生成(例如:给定一个pin,发现相关pins)或者在对候选排序时作为features使用。

出于便利性和泛化性,当我们描述PinSage算法时,我们会简单地将完整graph的node set使用 \(V = I \cup C\)来表示,不会显式对pin和board nodes进行显式区别(除非严格必要),会统一使用更通用的术语“node”.

3.2 模型结构

我们会使用localized convolutinal模块来生成nodes的embeddings。我们从input node features开始,接着开始学习nueral networks,它会将通过graph来对features进行transform和aggregate来计算node embeddings(图1)。

前向传播算法

我们会考虑生成一个embedding的任务, 对于一个node u的\(z_u\),它取决于node的input features和围绕该node的图结构。

图片名称

算法1

PinSage算法的核心是一个localized convolution操作,其中,我们会学习如何从u的邻居(图1)的信息聚合。该过程在算法1 CONVOLVE中有详解。基本思想是:我们将关于u的邻居通过一个dense neural network进行转换成representations \(z_v \forall v \in N(u)\),接着在vectors的结果集(第一行)上应用一个aggregator/pooling function(例如:一个element-wise mean或weighted sum,表示为\(\gamma\))。该aggregation step会提供一个关于u的局部邻居\(N(u)\)的vector representation \(nu\)。我们接着将aggretated neighborhood vector \(n_u\)与u的当前representation \(h_u\)进行拼接,并将contantenated vector \(n_u\)与u的当前representation \(h_u\)进行concatenate在一起,并通过另一个dense neural network layer(第2行)将concatenated vector进行转换。经验上,我们观察到:当使用concatenation operation来替代average operation时会取得极大的效果提升。另外,在第3行中的normalization会使得训练更稳定,它对于normalized embeddings来说执行近似最近邻搜索更高效(第3.5节)。该算法的output是一个关于u的representation,同时包含了它自己以及它的局部图邻居间的信息。

Importance-based neighborhoods

在我们的方案中,一个重要概念是:如何定义节点邻居\(N(u)\),例如:如何在算法1中选择要convolve的邻居集合。然而,之前的GCN方法会简单检查k-hop的图邻居,在PinSage中,我们定义了importance-based 邻居,其中:一个node u的邻居会被定义成T nodes会对node u施加最可能的影响。具体的,我们会从node u开始模拟random walks,并且计算出由random walk访问到的L1-normalized的节点访问数。u的邻居接着被定义成:对于node u具有最高的normalized visit counts的T个nodes。

该importnace-based邻居定义的优点是两面的。

  • 首先,选择一个固定数目的节点来聚合,允许我们控制着在训练期间该算法的memory footprint。
  • 第二,当聚合邻居的vector representations时,它会使算法1考虑邻居的importance

特别的,我们会在算法1中实现\(\gamma\)作为一个weighted-mean,它使用根据L1 normalized visit counts定义的weights。我们将该方法称为importance pooling

Stacking convolutions

每个时间,我们会使用CONVOLVE操作(算法1),我们会从一个node处获得一个新的representation,并且将多个这样的convolutions相互进行stack,以便获得围绕node u的局部图结构的更多信息。特别的,我们会使用多个关于convolutions的layers,其中:在layer k上的convolutions的inputs依赖于来自layer k-1(图1)的representations output,其中:intial(例如:“layer 0”)representations等于input node features。注意,在算法1(Q,q,W,w)中的模型参数会跨nodes共享,但在layers间的参数则不同。

算法2会详细说明:stacked convolutions 是如何为一个nodes的minibatch set M生成embeddings。我们首先计算每个node的邻居,接着使用K个convolutional迭代来生成关于target nodes的layer-K的representations。final convolutional layer的output接着通过一个fully-connected neural network进行feed来生成最终的output embeddings \(z_u, \forall u \in M\)

我们接着学习模型的full set参数:对于每个convolutional layer \((Q^{k}, q^{(k)}, W^{(k)}, w^{(k)}, \forall k \in \lbrace 1,\cdots,K \rbrace)\)的weight和bias参数,以及最终dense neural network layer \(G_1, G_2, g\)的参数。在算法1中的Line 1的输出唯度(例如:Q的column-space维度)会被设置为,在所有layers上均是m。出于简洁性,我们将所有convolutional layers的输出维度设置为相当(算法1第3行),接着我们将该size参数通过d进行表示。该模型的最终output维度(算法2)也被设置为d。

算法2

3.3 模型训练

我们会以一个监督学习的方式,使用max-margin ranking loss来训练PinSage。在该setup中,我们假设:我们已经访问了一个关于items L的labeled pairs集合,其中:在该set中的pairs \((q, i) \in L\)会被假设是相关的——例如:我们假设:如果 \((q, i) \in L\),那么item i是对于query item q的一个好推荐侯选。training阶段的目标是,最优化PinSage参数,以便在labeled set中的pairs \((q,i) \in L\)的output embedding 是非常接近的。

我们首先详细描述了margin-based loss function。根据此,我们给出了一个关于我们开发的多个技术的总览,它会导致PinSage的计算效率以及快速收敛,允许我们可以在数十亿节点的graphs和数十亿训练样本上进行训练。最终,我们描述了我们的curriculum-training scheme,它会提升推荐的总体质量。

Loss function

为了训练模型的参数,我们会使用一个Max-margin-based loss function。基本思想是:我们希望最大化正样本的inner product,例如:query item的embedding以及相应的相关item。同时,我们希望确认,负样本的inner product(例如:在query item的embedding和一个unrelated item间的inner product)要小于正样本由一些pre-defined margin给出的正样本。对于 \((z_q, z_i) \in L\)的node embeddings的单个pair的loss function:

\[J_G(z_q z_i) = E_{n_k \sim P_n(q)} max \lbrace 0, z_q \cdot z_{n_k} - z_q \cdot z_i + \Delta \rbrace\]

…(1)

其中:

  • \(P_n(q)\)表示了对于item q的负样本的分布
  • \(\Delta\)表示了margin hyper-parameter

我们会在下面解释负样本的采样。

使用large minibatches的Multi-GPU训练

为了在单机训练上充分利用多个GPUs,我们会以一个multi-tower的形式来运行forward和backward propagation。有了多个GPUs后,我们首先会将每个minibatch(图1底部)划分成等size的比例。每个GPU会获取minibatch的一部分,并使用参数的相同集合进行计算。在backward propagation后,在所有GPUs上每个参数的gradients会聚合到一起,会执行单频synchronous SGD。由于需要在非常大数目的样本上(数十亿规模)训练,我们会使用large batch sizes进行运营,范围从512到4096.

我们使用由Goyal[16]提出的相似技术来确保快速收敛,当处理大的batch sizes时,并保持训练和泛化的准确率。我们使用一个渐近的热启过程(gradual warmup produre),在第一个epoch,根据线性scaling rule从小到大增加learning rate。之后,learning rate会指数递减。

Producer-consumer minibatch构建

在训练期间,由于size很大,对于数十亿nodes的feature matrix和adjacency list会在CPU memory中放置。然而,在PinSage的CONVOLVE step,每个GPU过程需要访问在邻居中节点的邻居和feature信息。从CPU内容中访问来自GPU的数据并不高效。为了解决该问题,我使用一个re-indexing技术来创建一个sub-graph \(G' = (V', E')\)包含了nodes和它的邻居,它会在当前minibatch的计算中被涉及。一个小的feature matrix会只包含了与当前minibatch的计算相关的node features,会被抽取出来,以便顺序与在G’中的nodes的index相一致。G’的adjacency list和small feature matrix会在每个minibatch迭代中被feed到GPUs中,因此,在GPU和CPU间的通信在CONVOLVE step期间不需要,极大提升了GPU利用率。

训练过程会交替使用CPUs和GPUs。模型会在GPUs中计算,然而:抽取features、re-indexing、negative sampling会在CPUs中计算。除了使用multi-tower training的并行GPU计算之外,CPU计算会使用OpenMP,我们会设计一个 producer-consumer pattern来在当前迭代中运行GPU计算,在下一迭代上并行使用CPU计算。这会进一步减少一半的训练时间。

对negative items进行采样

negative sampling会在我们的loss function中会被使用,作为edge likelihood的normalization factor的一个近似。当使用大的batch sizes训练时会提升效率,我们抽样一个关于500 negative items的集合,会在每个minibatch中与所有训练样本进行共享。对比起对于每个node独立运行负样本来说,这会极大节约需要在每个training step期间被计算的embeddings的数目。经验上,我们不会观察到在两个sampling schemes间效果的一个不同之处。

在最简单的case中,我们会从items的整个集合中均匀抽样负样本。然而,确保正样本(items(q,i)的pair)的内积大于q,500 negative items的每个都太“easy”,对于要学习的系统来说,不会提供足够好的“resolution”。特别的,我们的推荐算法可以发现在20亿items的catalog间与q 最相关的1000个relevant items. 换句话说,我们的模型会在20亿items上能够区分/标识 1个item。但在500个random negaive items,模型的resolution只有1/500. 因此,如果我们从20亿items中抽样出500个随机负样本items,任何这些与query items更轻微相关的items的机会更小。因此,该学习大概率不会做出好的参数更新,不能区分非常相关items和轻微相关items。

为了解决上述问题,对于每个正的训练样本(例如:item pair (q,i)),我们会增加“hard” negative examples,例如:items一定程度上与query item q相关,但与正样本item i不相关。我们称这些为“hard negative items”。他们会根据query item q的Personlized PageRank score通过在graph中的items进行ranking来生成。在2000-5000上排序的items会被随机抽样作为hard negative items。如图2所示,hard negative examples对比起random negative examples与query更相似,对于模型排序来说挑战更大,强制模型去学习以细粒度的方式区分items。

图片名称

图2

通过training produre使用hard negative items,是需要训练至收敛的epochs数目的两倍。为了帮助收敛,我们开发了一个curriculum training scheme【4】。在训练的第一个epoch,不会使用hard negative items,以便算法快速发现在参数空间中的一个区域,其中:loss是相对小的。我们接着在后续epochs中添加hard negative items,聚焦于模型去学习如何区分高度相关pins和轻微相关pins。在训练的epoch n,我们会为每个item添加n-1个hard negative items到negative items集合。

3.4 通过MapReduce的Node embeddings

在模型被训练之后,直接使用训练好的模型来为所有items(包含了在训练期间没有见过的items)来生成embeddings仍然是很具挑战的。使用算法2以naive的方式计算nodes的embeddings,会导致重复计算,这是由nodes的K-hop邻居间的重合引起的。如图1所示,当为不同的target nodes生成embeddings时,许多nodes会在多个layers上重复计算。为了确保有效的inference,我们开发一个MapReduce方法,它无需重复计算即可运行model inference。

图片名称

图3

我们观察到:node embeddings的inference可以非常好地借助于MapReduce的计算模型。图3详述了在二部图上的data flow,pin-to-board Pinterest graph,其中,我们假设:input(例如:“layer-0”) nodes是pins/items(以及layer-1 nodes是boards/contexts)。MapReduce pipeline具有两个关键部分:

  • (1) 一个MapReduce 工作可以用于将所有pins投影到一个低维latent space中,其中:aggregation oepration会被执行(算法1,第一行)
  • (2) 另一个MapReduce job接着用于将产生的pin representations,与他们出现的boards的ids进行联合,board embedding的计算会通过它的邻居的features进行pooling。

注意:我们的方法会避免冗余计算,每个node的latent vector只会计算一次。在boards的embedding会被包含后,我们会使用两个多的MapReduce jobs来计算关于pins的第二层的embeddings,以一个相似的方式,该过程可以尽快迭代(直到K个convolutional layers)。

3.5 高效地最近邻lookups

由PinSage生成的embeddings可以被用于许多下游推荐任务中,在许多settings中,我们可以直接使用在学到的embedding空间上通过最近邻查询这些embeddings来做出推荐。也就是说,给定一个query item q,我们可以推荐那些关于query item的embedding的embedding是K个最近邻。ANN可以通过locality sensitive hashing有效获得。在hash function被计算后,items的检索可以使用一个two-level检索过程,基于Weak AND操作符来实现。假设PingSage model通过离线训练给出,所有node embeddings会通过MapReduce进行计算,并在database中保存,有效的最近邻lookup operation可以确保系统以一个在线的方式提供服务