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

摘要

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

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

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

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

• 输入数据：包含了稀疏的ID类型的Feature，比如：用户ID，物品类ID，以及稠密的non-ID Feature，如：图像，音频，视频等信息。
• Embedding层：每个稀疏的ID都有相对应的embedding表构成了embedding，其主要是将稀疏的ID feature通过embedding层得到embedding vector，将高维稀疏的向量信息，变成低维稠密的向量信息。
• Feature交互层：主要是将Embedding层输出的稠密向量信息，和输入数据部分的non-ID特征，交互到一起，形成一个包含完整输入信息的稠密特征；
• DNN层：由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数据给NNWorker，触发其发送一个request给EmbeddingWorker，查找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网络部分，每张卡保存完整的模型，数据则均分到每张卡中，相当于做了一个数据并行。

图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] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021

[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

 往期推荐

关于壁仞科技研究院

壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。

扫码关注我们