常用并行计算算子原理及其在LightGBM中的实现

一、概述

        在大规模机器学习中，需要应对巨大的训练数据及计算量。当单机遇到性能瓶颈时需要通过多台机器并行训练来弥补计算能力与内存的不足。采用并行方式进行机器学习时，常常分为模型并行与数据并行。模型并行是将模型拆分成多个分片，由几个计算节点分别持有，共同协作完成训练，适用于模型规模非常大的情形。数据并行是将数据拆分为不同的部分，分别存放在不同的计算节点上，同时每个计算节点都维护一个相同的模型。数据并行方式的训练过程中，每个计算节点只对自己存放的数据进行计算得到局部结果，然后再对局部结果进行全局归约。每个计算节点得到全局值之后，再进行模型更新。以下分别介绍并行计算中常用的归约算子，同时结合LightGBM的源码分析几个常用算子的实现方式。

二、并行计算中常用算子的主流实现方式剖析

        在并行计算中，数值归约是非常重要且使用频率非常高的一种操作。数值归约算法实现方式的好坏直接影响整个计算过程的速度。常用的一些归约算子包括：Allreduce、Reduce、Allgather、Broadcast、ReduceScatter等。下面分别对这几种算子进行介绍。

        一些定义：N－计算节点的总数； R_i － 第i个计算节点； V_i － 第i个计算节点的局部数值； root － 树状网络的根节点；k － 2^k <= N的最大整数； remain = N - 2^k； reducer － 归约操作；

2.1 Allgather

        Allgather操作的目的是将计算节点的本地数据V_i同步到其他所有的计算节点，使得每一个节点都拥有一份V₀ - V_N-1的值。主流的实现方式有Recursive Doubling、Bruck Algorithm。

2.1.1 Recursive Doubling

图1 Recursive Doubling算法的Allgather过程

上图描述了Recursive Doubling算法的工作过程。

第一步：在每一个节点上开辟sizeof(V₀) + sizeof(V₁) + sizeof(V₂) + sizeof(V₃)的内存，并且将本地的数据拷备到对应的起始位置；

第二步：如图所示，step-1时，节点与其相距2⁰的节点进行数据互发；

第三步：step-2时，节点与其相距2¹的节点进行数据互发，发送的数据包括节点自身的数据以及之前接收的数据；

依此方式执行k步，最终所有的计算节点都拥有一份全局的数据了。

Recursive Doubling方式要求N为2的整数倍，计算节点收集到所有数据的通信次数为log(N)。

2.1.2 Bruck Algorithm

图2 Bruck 算法的Allgather过程

上图描述了Bruck Algorithm算法的工作过程。

第一步：如图，R₁将数据发给R₀, ......., R₀将数据发给R₄，发送一个数据块；

第二步：重复k步，第k步时，R_i将数据发送给R_i-2^k，发送个数据块min(2^k, N-sum_block)；sum_block为计算节点上已有数据块大小。

第三步：调整数据块的顺序，调整方法在4.1节代码中介绍；

以上这种方式不要求N一定为2的整数倍，计算节点收集到所有数据的通信次数为ceil(log(N))。

2.2 Broadcast

        Broadcast操作的目的是将主节点的数值广播到其他所有的计算节点。主流的实现方式有Binomial Tree Algorithm、Geijin Algorithm。

2.2.1 Binomial Tree Algorithm

图3 Binomial Tree 算法的Broadcast过程

上图描述了Binomial Tree Algorithm算法的工作过程，一般需要做broadcast操作时都有一个主节点设为R₀。各个计算节点的网络连接如上图。

第一步：R₀将本地值发送给左右孩子节点R₁和R₂，同时将R₁和R₂设置为root，再往它们的左右孩子节点发送；

第二步：第i步，R_i将本地值发送给R_2*i+1和R_2*i+2，同时将R_2*i+1和R_2*i+2设为root，再重复这一过程，直到该节点没有孩子节点；

        Binomial Tree Algorithm这种广播方式，每一次发送的数据块为整个数据块的大小，在数据块较小时比较适用。但当数据块较大时，以下介绍的Geijin Algorithm会更加适用。

2.2.2 Geijin Algorithm

Geijin Algorithm适用于大的数据块。主要思想是：

第一步：将需要Broadcast的数据分成N块，同时Scatter到各个计算节点上；

第二步：对各个计算节点分发到的子数据块执行Allgather，即可完成操作；

Geijin Algorithm相比于Binomial Tree Algorithm增加了通信次数，但是每次收发的数据块小了，对于大数据块更能降低带宽消耗。                

2.3 Gather

图4 Binomial Tree 算法的Gather过程

Gather操作的目的是将所有计算节点上的数据集合到主节点上。也可以采用Binomial Tree Algorithm算法。各个计算节点的网络连接如上图。

第一步：从叶子节点开始，将本地数据发送到其父节点；

第二步：第一步中的父节点接收到数据后，继续将本地数据以及接收到的数据往上发，最终根节点上的数据就是所有节点数据Gather之后的结果了；

Binomial Tree Algorithm的实现方式通信次数为log(N)。

2.4 Reduce

        Reduce操作的目的是对所有计算节点上的数值进行归约操作，并将归约后的结果保存到主节点上。主流实现方式有Binomial Tree Algorithm以及Rabenseifner提出的算法[2]。

2.4.1 Binomial Tree Algorithm

        采用二叉树的方式，类似于Broadcast过程的逆过程。首先需要对所有的计算节点按照二叉树建立网络连接。然后从叶子开始，将本地数据发送到父节点，父节点在接收到数据后，按照给定的归约函数执行一次归约操作，再将归约后的结果发送到其父节点，重复这一过程直到根节点。最终根节点上的值就是Reduce之后的结果。

2.4.2 Rabenseifner's Algorithm

        该算法适用于数据块较大的情形，首先进行一次ReduceScatter（2.5介绍）操作，然后再进行一次Gather（2.3）操作。

2.5 Allreduce

        Allreduce操作的目的是对所有计算节点上的数值进行归约操作，同时每一个计算节点均获得归约后的结果。主流实现方式有Binomial Tree Algorithm以及Recursive Doubling。

2.5.1 Binomial Tree Algorithm

        如果采用二叉树算法，可以通过2.4.1介绍的Reduce + 2.2.1介绍的Broadcast两次操作完成。如下图所示。

图5 Binomial Tree 算法的Allreduce过程

        这种方式包括push up和pull down两个过程，push up是将本地数值进行上发，同时在接收端进行归约操作并且将得到归约后的结果继续上发。pull down是在root节点上得到了全局归约值之后的一个Broadcast过程。

2.5.2 Recursive Doubling

图6 Recursive Doubling算法的Allreduce过程

上图中描述了使用Recursive Doubling算法进行Allreduce操作的过程。具体步骤如下：

第一步：首先对N个计算节点两两分组，如R₀与R_0+2⁰一组。每组之间的计算节点相互发送数据，接收方将数据存放到缓存区(step1 → step2)；

第二步：每一个计算节点，对缓存区数据与本地数据做一次归约函数操作reducer(V_i, V_recived)(step2 → step3)；

第三步：再重新两两分组，如R₀与R_0+2¹一组。每组之间的计算节点相互发送数据，接收方将数据存放到缓存区(step3 → step4)；

第四步：依此重复执行k次之后，每个节点上的结果就是最终的归约结果；

这种方式的Allreduce要求N为2的整数倍，完成整个操作需要的的通信次数为log(N)。当N非2的整数倍时，可以采用2.6.2的方式执行一个辅助步骤再进行。

2.6 ReduceScatter

        ReduceScatter操作的目的也是对所有计算节点上的数值进行归约操作，但是各个节点只保留归约后的部分结果。以下介绍Recursive Halving的实现方式。

2.6.1 当workers数量为2的k次方时，以下以8个workers为例进行介绍：

        最初，每一个worker都将数据分成8个数据块，每个数据块用Fi表示。本例中ReduceScatter的目的是使R₀得到reducer(0_F0, ...., 7_F0)的结果、R₁得到reducer(0_F1, ...., 7_F1)....。

图7 Recursive Halving算法的ReduceScatter过程

上图中描述了使用Recursive Halving算法进行ReduceScatter操作的过程。以R₀为例具体步骤如下：

第一步：先计算每一个节点需要获得数据块哪一部分的归约结果，并将数据块切分成N个子块；

第二步：第一次通信，R₀将4-7子块发送给R₄，并接收R₄发送过来的0-3子块，然后在本地对自身的数据以及接收到的数据进行归约操作；

第三步：第二次通信，R₀将2-3子块发送给R₂，并接收R₂发送过来的0-1子块，然后在本地对自身的数据以及接收到的数据进行归约操作；

第四步：依次执行，最终每个节点在对应的位置都得到了分配给其的结果，如图中红色位置；

这种方式的ReduceScatter要求N为2的整数倍，完成整个操作需要的的通信次数为log(N)。

2.6.2 当workers数量非2的k次方时，可以采用以下方法：

图8 workers数量非2的k次方时Recursive Halving的辅助步骤

第一步：小于2*remain的节点中，R_i为偶数的节点，将它需要做ReduceScatter的所有数据发送给R_i+1；

第二步：设置一个virtual_rank(VR)。VR的计算：第一步中偶数计算节点设置为-1，奇数的设置为0开始依次递增。大于2*remain的计算节点VR=R-remain。详细做法见4.2代码；

第三步：virtual_rank非-1的节点按照2.6.1的Recursive Halving方法进行操作即可；

这种方法相比于2的k次方时多了第一步的所有数据的发送步骤。

三、LightGBM网络组网步骤

        LightGBM并行计算中，计算节点之间采用静态互连网络，即程序执行期间，节点与节点之间的连接保持不变。LightGBM网络的建立需要事先定义所有worker的ip以及port，组网的过程通过Network::Init()完成，主要步骤如下：

3.1 网络初始化入口

        

        建立网络连接主要的逻辑在new Linkers(config)中。 

3.2 LightGBM中网络初始化过程

        初始化过程主要是建立Linkers对象：

        网络初始化过程中，首先从配置中读取当前worker的端口号。然后解析machine list文件，获取所有worker的ip列表以及port列表。接下来需要对Bruck算法及Recursive Halving算法计算出当前worker需要连接到其他哪些worker。

        

        从2.12可以看出BruckMap不要求workers数量为2的k次方，因此对于任意数量的workers，都可以按以上步骤计算当前节点需要在第k步通信时发送数据给哪个worker以及接收哪个worker发送过来的数据。

        从2.6.1可知，Recursive Halving算法要求workers数量为2^k，但当workers数量非2^k时，可以通过2.6.2介绍的方法进行解决这一问题。以上construct步骤主要是计算哪些worker需要参于到Recursive Halving过程中，同时计算第k次通信配对的worker，以及第k次通信发送、接收的数据块大小及位置（此处的值只是一个基数，具体位置及大小在实际的执行过程中根据实际数据块的大小计算）。

 

       计算完需要连接的worker之后，Construct将依次连接到需要连接的worker上，并等待被连接的worker的连接。连接与监听过程分别在两个线程中执行。自此，LightGBM静态互连网络的建立已经完成。

四、LightGBM中Allreduce、ReduceScatter、Allgather实现解析

        LightGBM的并行训练过程中使用了大量的数值归约操作，如随机种子的同步用到了Allreduce、对数据进行分桶时计算bin_mapper用到了Allgather、计算最佳分裂时用到了ReduceScatter等，以下分别对这三种归约操作的代码进行解析。

4.1 Allgather代码解析

LightGBM中Allgather的实现采用Bruck Algorithm算法：

        以上的操作经过log(N)次的通信能够完成Allgather。每一次通信数据块的大小为2^k*send_size。可以理解，这种方式适用于N不太大并且数据块大小不太大的情形。

4.2 ReduceScatter代码解析

        以上仅介绍非2^k个节点的情形，workers为2^k时执行过程仅需要第二个for循环。

4.3 Allreduce代码解析

        Allreduce根据数据块的大小的不同，分别采用两种方式进行操作。在需要归约的数据块较小时采用Allgather加上一个reducer操作完成。针对数据块较大的情形则采用另一种方式：ReduceScatter + Allgather操作。

4.4 一个调用Allreduce的例子

以下介绍一个使用Allreduce的例子，在训练之前以label的平均值作为初始值：

        因为init_score的size小，因此先通过Allgather收集所有worker上的init_score，然后在本地进行reducer操作，此处的reducer是做sum。

五、总结

        本文主要是通过对MPI中主流归约算子的理解，进行归纳并且通过实例化的形式进行总结。然后结合LightGBM的源码，总结了一般并行计算中各个节点进行网络连接的过程，以及LightGBM对Allreduce、Allgather、ReduceScatter三种数值归约方式的代码实现。一般并行计算中的归约算子很多，并且每个算子都有多种实现方式，各种不同的实现适用于不同的场景，有些在小数据块时效率高，有些在大数据块时更有优势，如何进行选择需要根据业务场景甚至通过实验对比来确定。

六、参考

1. Generalisation of Recursive Doubling for AllReduce

2. Optimization of Collective Communication Operations in MPICH

3. https://github.com/Microsoft/LightGBM

4. https://www.zhihu.com/question/37057945