飞桨推出异构参数服务器架构，异构硬件高效组合，训练速度提升65%以上- DoNews

“双11”了，对于广大网购爱好者来说那绝对是不可错过的狂欢时刻！当今网购之所以如此火爆，不仅仅是营销策划的作用，智能化的搜索推荐技术也可以说是功不可没。它能把你日思夜想或者潜意识中动过购买念头的商品通通推送到你的面前，甚至会让人有一种冥冥自有天意、不买对不起上苍的感觉。而这背后往往都会有深度学习领域中个性化推荐模型发挥着威力。为了能够更准确的预知用户的内心需求，快速训练出效果良好的推荐模型并尽快部署上线，成为了各大网购业务相关企业的共同追求。

在搜索推荐领域中，通常推荐模型的网络并不复杂，但是由于对应的特征空间大，所以通常会使用算力要求不高但是规模非常庞大的Embedding和FC层来将用户及商品的高维稀疏特征向量转化为低维的稠密特征向量，导致这类模型的参数维度可以达到千亿甚至万亿级别，并且还伴随着大规模稀疏的特点。同时各个企业为了模型效果，往往会使用尽可能多的数据来训练模型，尤其是其中的点击数据，如果来自一些带有“选择困难症”的人，可能会非常的长，导致数据规模愈加庞大，甚至可以达到亿级。因此，相比CV和NLP领域而言，在搜索推荐场景中，单次Batch训练中前后向计算的时间远低于数据读取、拉取和更新参数等过程的时间，再加上庞大的数据和参数量，导致内存要求很高，甚至可以达到几十T，而拥有这类特点的训练任务通常都被称为IO密集型训练任务。

那么如何快速完成IO密集型训练任务呢？

单机训练肯定无法承受这样的任务，只有分布式训练才是解决问题的方法。简单来说，分布式训练就是使用多个硬件设备共同完成训练任务，可以分为All-ReduceCollective和ParameterServer(参数服务器)两种训练架构。

在Collective架构中，集群中存在多个地位平等的Trainer（也可以叫Worker）。在数据并行的方式下，每个Trainer上都保存一份完整的模型网络结构和模型参数，在前向和反向计算过程中，每个Trainer使用自己划分的数据（shard）进行计算，得到对应的梯度；然后Trainer之间通过All-Reduce等通信方式同步梯度到所有Trainer，最后每个Trainer使用同步后的梯度独立完成参数更新。

图1Collective架构

由于推荐搜索领域模型的Embedding层规模庞大以及训练数据样本长度不固定等原因，导致容易出现显存不足和卡间同步时间耗费等问题，所以Collective架构很少被用于搜索推荐领域。

再来看看参数服务器架构。与Collective相比，参数服务器架构中除了Trainer之外还有一个角色叫Server。在执行前向与反向计算过程中，Server负责从各个Trainer收集汇总梯度，然后使用这些梯度信息计算更新参数，并向各个Trainer分发更新后的参数信息。

图2传统参数服务器架构

但是在应对IO密集型任务时，传统参数服务器模式会出现以下问题：

如果使用的是CPU机器作Trainer，通常CPU核心数较多，且购置价格便宜，训练中可以充分利用CPU多核的优势，使用异步训练模式，在简单模型上可以极大提升数据吞吐量，对提升整体训练速度有不俗的表现。但是AI开发者为了追求模型效果开始逐步在推荐模型中增加复杂网络部分，CPU计算能力的弱势便暴露无遗，训练耗时会变得不可接受。

图3传统参数服务器架构（CPU机器）遇到算力瓶颈

如果更新集群硬件，改为使用的是GPU机器作Trainer，则可能会出现资源利用率低和网络带宽不足的问题：

资源利用率低：IO密集型任务主要还是数据读取和模型读取(IO)的性能对训练速度更具决定性。在这方面使用GPU的集群如同“杀鸡用牛刀”，并不会对整体的训练速度产生多少帮助，而且硬件资源得不到充分利用。而如果加大batchsize提高单次前向后向计算占比，也会受到单机IO吞吐能力的制约，而且由于模型高度稀疏，模型效果也会受到影响。

网络带宽不足：由于一台GPU机器可以顶替多台CPU机器，必然会导致整个集群网卡数量降低。这样在将数据和模型下载到Trainer阶段就很容易出现带宽瓶颈。尤其是当Server和Trainer都在一台设备时，一旦每个Trainer划分的训练数据过大，训练准备阶段就会变得极为耗时，整个训练过程直接输在了起跑线上。

图4传统参数服务器架构（GPU机器）遇到IO瓶颈

也许有人会提出直接将IO任务交给GPU机器上的CPU不就可以解决了吗？答案是否定的。一台GPU机器上的GPU和CPU的硬件配比是固定的，且单台GPU的多卡相比多台CPU机器而言，每个GPU卡对应的CPU核数相对较少，这就导致GPU前向后向训练的越快，对CPU读数据和模型参数的要求就越高，这样CPU反而更容易成为瓶颈，而且不能解决网络带宽不足的问题。

新型算力接入成本较大。随着AI芯片发展日新月异，各种高算力低成本的芯片已进入工业实用化阶段。在参数服务器模式下，一旦更换新型算力硬件，需要完成计算集群的迁移，软件栈变更，训练速度及效果打平等一系列工作。

那么有没有可以高效调整机器配比，快速支持新硬件的接入，使用GPU等高算力的同时还不用担心IO瓶颈的方法呢？

鱼与熊掌可以兼得？——异构参数服务器训练架构

上面的要求好像很多，其实核心问题就在于硬件的配置，传统的参数服务器对硬件的统一性要求太严格，而现实是单一“兵种”是无法应对大部分“战场”的。如果能同时使用不同的硬件同时参与到训练中，让它们各司其职，发挥各自优势，很多问题就可以迎刃而解了！

为了应对上述问题，飞桨框架2.0版本基于工业实践，创新性地推出了大规模稀疏异构参数服务器功能，一举解除了传统参数服务器模式必须严格使用同一种硬件型号Trainer节点的枷锁，使训练任务对硬件型号不敏感，即可以同时使用不同的硬件进行混合异构训练，如CPU、AI专用芯片（如百度昆仑XPU）以及不同型号的GPU如V100、P40、K40等。同时还可以解决大规模稀疏特征模型训练场景下，IO占比过高导致的芯片资源利用率过低的问题。通过异构参数服务器训练架构，用户可以在硬件异构集群中部署分布式训练任务，例如云服务器集群，实现对不同算力的芯片高效利用，为用户提供更高吞吐、更低资源消耗的训练能力。

图5异构参数服务器架构示意图

异构参数服务器架构的原理

如之前所述，在传统参数服务器模式下，前向及反向步骤在Trainer端完成，参数更新在Server端完成。而在异构参数服务器架构中，原先那样粗粒度的任务分工显然不再适用，只有让不同硬件参加到训练任务中，并在自己最擅长的领域工作，才能各显神通。因此飞桨的研发人员在传统的CPU-Trainer之外，又部署了一类Heter-Trainer，并在Heter-Trainer上配置有Service。这样CPU-Trainer便将Heter-Trainer视为特殊的Server，从而对Heter-Trainer的异构算力不敏感，使Heter-Trainer可以混部各类不同型号芯片（例如GPU或XPU等）。

为了让CPU-Trainer和Heter-Trainer完成不同类型的任务，飞桨将整个计算任务做了进一步的拆分。CPU-Trainer继续负责数据读取、参数的读取和更新、参数的网络传输等IO任务，甚至在小BatchSize训练时，使用CPU处理更快的网络层计算；而对于复杂网络的计算，则CPU-Trainer会交给Heter-Trainer处理。总之，所有CPU相对GPU或XPU更擅长处理的操作都放在CPU中，其它的放在GPU或XPU中。

图6传统参数服务器架构的异构改造

值得注意的是，异构参数服务器架构并不只是简单的将计算任务拆分。拆分后，异构硬件间的通信代价高的问题也需要解决。为了解决这个问题，飞桨的整个传输通信过程得到了优化。该过程将不会完全依赖于硬件自身的吞吐，而是引入类协程任务调度机制，该机制采用了异步等待和多Queue流水线等方式对任务进行灵活调度，避免传输队列阻塞，对任务进行灵活调度，并且还优化了Heter-Trainer的吞吐机制，高效的避免传输队列阻塞。

此外，异构参数服务器训练架构在各个硬件设备上还引入了异构存储机制，将模型训练所用到的参数，按照一定层级分别存储到SSD、内存和显存中，并且支持参数在SSD、内存、显存之间的快速拷贝，使模型规模不再受限于内存或显存的规格，可以进一步支持更大参数规模的模型训练。

如果用户希望将现有GPU机器上的传统参数服务器架构改造为异构参数服务器架构，则用户可以通过飞桨从机器上划分出部分CPU资源形成CPU-Trainer，然后再将原有的GPU-Trainer改造为Heter-Trainer。但是通常在此场景下，CPU-Trainer所能处理的数据，仅需要使用少量GPU资源组成的Heter-Trainer处理，便足以完成训练计算任务，从而可以节省大量的GPU资源。

而对于将CPU机器上的传统参数服务器架构改造为异构参数服务器架构的需求，飞桨可以支持用户根据任务所需的实际算力灵活引入异构硬件（GPU或XPU）成为Heter-Trainer，这样不仅可以保障资源的利用率，而且可以利用用户手中的旧型号硬件机器，大大节约用户的成本。这一特点也使异构参数服务器架构非常适合部署在云上异构集群场景中。

图6传统参数服务器架构的异构改造（外挂硬件，CPU机器为例）

飞桨的异构参数服务器架构不仅可以支持CPU机器、GPU机器和XPU机器的混合模式，而且支持全部CPU-Trainer或GPU-Trainer的那样的传统参数服务器的硬件组合，甚至在GPU单类型机器的集群内采用Heter模式，可以达到训练提速的效果。

异构参数服务器的优势

异构参数服务器兼顾了传统参数服务器架构的大规模稀疏及异步优势，充分利用了GPU等AI芯片带来的算力上的提升，在模型训练速度上显著提升。

我们基于飞桨和其它框架，在如下架构上进行测试：

使用4台GPU机器作Trainer的Collective架构（Collective-GPU）。

使用4台CPU机器作Trainer的非异构参数服务器架构（PS-CPU）。

使用4台GPU机器作Trainer的非异构参数服务器架构（PS-GPU）。

使用4台CPU机器和若干GPU机器作Trainer的异构参数服务器架构（PS-Heter）。

其中CPU机器型号为Intel(R)Xeon(R)Gold6148，GPU机器型号为Intel(R)Xeon(R)Gold6148+TeslaV10032G。

首先我们在简单的CTR-DNN模型（FC层维度分别为“400->400->400->2”）和Criteo数据集进行测试，模拟简单IO密集型训练任务。如下图所示，我们可以看到对于简单模型，飞桨异构参数服务器架构（PS-Heter，4CPU+1GPU）性能表现优于其它训练架构。

当我们扩大FC层维度（FC层维度分别为“4096->2048->1024->2”），模拟复杂网络时的IO密集型训练任务，如下图所示，可以看到PS-CPU架构的性能下降显著。