AllReduce

背景

深度学习领域的关键挑战之一，就是耗时。耗时的原因来自于两个方面，一方面是模型越来越大，参数越来越多，千万甚至数十亿，另一个方面是训练数据量的增多。所以更低的耗时、更快速的实验，对于模型的的开发效率至关重要，使研究者有机会探索更多的超差组合。

分布式深度学习算法是降低耗时问题的重要方法。2017年6月，Facebook发布论文，文中说他们用32台服务器256块GPU，将Resnet-50模型在ImageNet数据集上的训练时间从两周缩短到1小时。他们采用的就是分布式学习算法。

在分布式训练中，数据并发是经常采用的架构，这种架构下每个节点都有完整的模型结构，共同消费训练数据，达到加速训练的目的：

学习的步骤为：

1. 每个节点分别消费minibatch， 计算各自的梯度
2. reduce所有梯度，一般reduce操作采用求平均值
3. 梯度的reduce结果用来更新模型参数
4. 回到步骤1，直到数据消耗完成

这里问题就出现了，就是reduce怎么实现呢 ？

首先我们把这个问题抽象成一个一般性的问题.

问题定义

存在N个节点，$P_i$表示第$i$个节点，$i\in [0,N-1]$。每个节点都有一个数组$A_i$，数组长度也为N, 即 $len（A_i）=N$。求目标数组R，满足以下条件：

$$R[k]=f(A_0[k],A_1[k],...,A_{N-1}[k])$$
$$len(R)=N$$

其中reduce函数 $f$ 可能是平均值函数、求和函数、最大值函数、最小值函数，等等。

比如，具体对应到上面GPU数据并行训练的场景，节点$P_i$表示每个GPU，数组$A_i$表示每个GPU各自计算的梯度，数组$R$表示梯度的聚合结果:

定义完问题之后，我们来讨论下这个问题的算法设计。一个常见的算法思路，是所有GPU将各自的数据发往一个中心节点，中心节点让后将reduce的结果返回给GPU，如下所示:

这个算法的缺点主要有两点：

• 中心节点需要等待所有其他节点的数据
• 中心节点的通信量比较大，并且通行量正比与节点的数量

Ring-AllReduce

下面用四个节点的例子，说明ring-allreduce的过程。Ring-AllReduce算法分为两个阶段。

第一阶段, scatter-reduce阶段：

step 1:

step 2:

step 3:

step 4:

step 5:

step 6:

step 7:

第二阶段，all-gather阶段：

step 1:

step 2:

step 3:

step 4:

step 5:

Tree-AllReduce

Tree-AllReduce算法过程也分为两个阶段:reduce与broadcast.