推荐场景训练加速:大规模混合分布式系统

发布时间:2022-04-25 07:00

摘要

 

 

当前深度学习模型,已经被广泛采用在推荐系统领域中。为了获取更好的精度,模型以及数据集的规模也越来越大,这些使得分布式训练系统面临着越来越大的挑战。而这其中最大的挑战来自于处理非常大非均匀的训练任务:一方面,模型的嵌入层包含了大部分的模型参数,训练时其运算特征是访存密集型;另一方面,模型的神经网络层,其运算特征则主要是计算密集型。为了解决这一类不均匀任务带来的挑战,一些混合式的分布式训练系统被提出,本文主要对这类混合系统做简要的分析。

 

推荐场景深度学习模型的介绍

 

 

推荐系统是互联网服务中的一个非常重要的部分,其任务包含:CTR(Click-Through Rate),BTR(Buy-Through Rate)的预测,被广泛应用在广告,视频和购物等领域中。随着深度学习模型在各个领域的广泛应用,推荐场景也引入了深度学习模型,例如:Wide Deep Learning,DCN,DIEN等,均取得了非常好的效果。

一个典型的推荐场景的深度学习模型的结构,如下图1所示,模型主要包含了处理稀疏和稠密两部分特征的不同网络结构。整个模型的主要内容包括:

l输入数据:包含了稀疏的ID类型的Feature,比如:用户ID,物品类ID,以及稠密的non-ID Feature,如:图像,音频,视频等信息。

lEmbedding层:每个稀疏的ID都有相对应的embedding表构成了embedding,其主要是将稀疏的ID feature通过embedding层得到embedding vector,将高维稀疏的向量信息,变成低维稠密的向量信息。

lFeature交互层:主要是将Embedding层输出的稠密向量信息,和输入数据部分的non-ID特征,交互到一起,形成一个包含完整输入信息的稠密特征;

lDNN层:由ID信息和non-ID信息交互的稠密特征,通过相应的DNN网络,提取得到相应所需要的预测信息。

 

图1 推荐系统深度学习模型结构示意图[1]

   

整个系统中,模型的参数主要包含两部分:Embedding层的嵌入向量,以及DNN层的模型参数。其中Embedding层的参数量占了模型参数的绝大部分,该部分在训练过程中,会产生大量的查表访问分布式存储的行为,因而是访存密集型,而DNN层的参数量相对较少,但计算量比较大,因而是计算密集型。如何协调好访存和计算,是影响设计系统性能的关键。

 

 

混合分布式系统

 

 

随着推荐系统中深度学习网络模型的发展,其模型的规模也变得越来越大,甚至达到100 Trillion量级,一个分布式系统在如此级别,通常需要包含成百上千的机器构成的机群,而这也使得存储和计算不均性的问题,愈发严重。为了解决这些问题,一些混合的分布式系统被提出,主要是将稀疏特征部分以及稠密特征部分的系统分别独立设计,采用不同的并行及同步方案,构成了一个混合的系统。

稀疏特征部分的系统,主要解决访存密集型的问题,这一部分由于查表过程的分散性,主要采用PS(Parameter Server)的分布式方式进行设计。而由于embedding层参数量巨大,为了提高查表的效率,多层次的存储以充分利用各类资源的IO访问带宽,梯度异步同步等方式被采用。

稠密特征部分的系统,其DNN模型的参数量一般不是特别大,基本上沿用DNN所用的All-Reduce分布式同步方案,其并行的方式也可以采用常用的如数据并行进行分布式加速,这一部分一般直接复用框架所提供的分布式方案。

混合系统的设计,会有不同的方案呈现,本文主要介绍两个系统:Persia和HugeCTR两个系统,来分析设计的原理。

 

Persia系统

 

Persia系统[1]是由快手开发的一个加速CTR的分布式混合系统。如图2所示,系统主要包含四个部分:Embedding PS,Embedding Worker,NN Worker,Data Loader。Embedding PS主要负责Embedding层的参数存储,以及前向查表和反向embedding梯度的更新;Embedding Worker主要负责从Data Loader中获取sample数据,并收集多个sample数据合并成batch,以及处理稀疏数据与Embedding PS之间查表请求发送;NN Worker主要负责DNN网络的部署,包括接收稠密输入信息,并对网络进行前向与反向计算。其完整工作流包括:

 

图2 Persia系统的结构及流程示意图[1]

 

1.Data Loader发送sample数据给Embedding Worker,同时其会分配一个index(ε)给该sample,并将index返回给Data loader;

2.Data Loader发送index所对应的Non-ID数据给NN Worker,触发其发送一个request给Embedding Worker,查找Index所对应的Feature ID,并进一步触发Embedding PS查找相应的embedding vector,将其发送给NN Worker;

3.NN Worker收集完整的sample数据后,并持续收集达到batch数量时,会触发FC(Forward Compute)和BC(Backward Compute)的计算,并产生相应的embedding gradient;

4.得到的embedding gradient,会在Embedding Worker中找到相应的Feature ID,并发送给PS进行参数更新,从而完成整个训练。

整体而言,Persia将稀疏网络处理的这一部分放到了CPU上,而dense部分则放到了GPU中,通过CPU+GPU的方案实现混合的分布式系统。

除了设计CPU+GPU分别处理稀疏网络和DNN网络的混合系统外,Persia对训练的流程的异步化也做了非常有意义的工作。如图3所示,完全的同步的流水线,其步骤依次包括:获取嵌入向量,前向计算,后向计算,Dense梯度同步,Embedding梯度更新。而这其中的异步处理包括:嵌入向量和前向计算可以异步处理;以及后向计算与Dense梯度同步也可以异步处理;另外由于Embedding向量的稀疏性,异步梯度更新的参数过时问题并不严重,因而embedding梯度更新与查表也可以做成异步处理。通过这些异步操作,可以极大地缩短流水线长度,提高训练的迭代速度。

 

图3  Persia异步化的处理流程示意图[1]

 

 

HugeCTR系统

 

 

HugeCTR[2]是英伟达开发的一个基于GPU加速的专用分布式训练框架设计,可以支持WDL (Wide and Deep Learning),DCN (Deep Cross Network) 等常用的CTR网络部署。

HugeCTR的主要特点,是在训练中将一个pass所需要的embedding表项,作为一个子表,分布式地存储在GPU的显存中,这样在batch训练过程中,查表不再在CPU侧的PS(parameter Server)中,而是在不同的GPU中查找,通过GPU的高速片内及片外IO带宽,可以极大提高查表效率。如下图4所示,HugeCTR的GPU主要架构,主要包含了两部分:1. 每次训练的一个迭代周期内,所需要的Embedding子表分布式地存储在每张GPU卡显存中,相当于是做了一个模型并行;2. NN网络部分,每张卡保存完整的模型,数据则均分到每张卡中,相当于做了一个数据并行。

 

https://nvidia-merlin.github.io/HugeCTR/master/_images/fig1_hugectr_arch.png

图4 HugeCTR架构示意图[3]

 

HugeCTR的加速原理如下图5所示,其主要的步骤包括:

1.每张卡GPU根据每次训练迭代周期(Pass)所需要的embedding 参数,从CPU侧的PS(包含完整的embedding表)拉取相应的表项构成子表;

2.训练过程中Sparse特征的查表,前向计算,反向计算,反向dense梯度更新,以及反向embedding梯度更新,都发生在GPU中;

3.完成一次迭代周期(Pass)后,再将GPU的子表表项上传并更新到CPU侧的PS中,然后进入下一次的迭代训练中。

整个过程中,除了第一步从CPU中拉取部分表项建立子表,以及最后一步将一个迭代周期(Pass)的embedding参数梯度更新到CPU PS总表中外,其余的查表及计算行为都发生在GPU中,而GPU在通信、存储和计算方面的性能优势,也可以大大加速整个训练迭代流程,从而可以获得更好的性能。

 

图5  HugeCTR加速方案原理图

 

HugeCTR与Persia不同的是,HugeCTR的Sparse网络处理以及DNN网络处理都部署在GPU中,这种方案也被称为全GPU的方案。

 

 

结论与思考

 

 

本文介绍了推荐场景中的混合分布式训练系统的内容,介绍了推荐场景中的深度学习网络模型的结构,以及训练过程中的挑战,并以此介绍了混合分布式系统如何应对这些挑战,提高计算效率,并以Persia和HugeCTR为例子详细介绍了混合分布式系统的实现。由于水平有限,文中存在不足的地方请各位读者批评指正,也欢迎大家一起参与我们的讨论。

 

 

 

参考文献

 

 

1.Xiangru Lian, etc. PERSIA: An Open, Hybrid System Scaling Deep Learning-based Recommenders up to 100 Trillion Parameters.

2.Even Oldridge, etc. Merlin: A GPU Accelerated Recommendation Framework, https://irsworkshop.github.io/2020/publications/paper_21_Oldridge_Merlin.pdf. 2020.

3.https://nvidia-merlin.github.io/HugeCTR/master/hugectr_user_guide.html

 

上一个: Kubric:高效地合成视觉数据集

下一个: 比大更大:Pathways上实现的大语言模型PaLM

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多