baidu的aibox介绍

Reading time ~1 minute

介绍

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.实验

参考

google Titans介绍

google在《Titans: Learning to Memorize at Test Time》提出了区别于Transformer的的一种新架构:Titans。我们来看一下它的实现,是否有前景:# 摘要在过去的十多年里,关于如何有效利用循环模型(recurrent mo...… Continue reading

meta QuickUpdate介绍

Published on January 02, 2025

kuaishou CREAD介绍

Published on August 05, 2024