背景
深度学习领域的关键挑战之一,就是耗时。耗时的原因来自于两个方面,一方面是模型越来越大,参数越来越多,千万甚至数十亿,另一个方面是训练数据量的增多。所以更低的耗时、更快速的实验,对于模型的的开发效率至关重要,使研究者有机会探索更多的超差组合。
分布式深度学习算法是降低耗时问题的重要方法。2017年6月,Facebook发布论文,文中说他们用32台服务器256块GPU,将Resnet-50模型在ImageNet数据集上的训练时间从两周缩短到1小时。他们采用的就是分布式学习算法。
在分布式训练中,数据并发是经常采用的架构,这种架构下每个节点都有完整的模型结构,共同消费训练数据,达到加速训练的目的:
学习的步骤为:
- 每个节点分别消费minibatch, 计算各自的梯度
- reduce所有梯度,一般reduce操作采用求平均值
- 梯度的reduce结果用来更新模型参数
- 回到步骤1,直到数据消耗完成
这里问题就出现了,就是reduce怎么实现呢 ?
首先我们把这个问题抽象成一个一般性的问题.
问题定义
存在N个节点, 表示第 个节点, 。每个节点都有一个数组 ,数组长度也为N, 即 。求目标数组R,满足以下条件:
其中reduce函数 可能是平均值函数、求和函数、最大值函数、最小值函数,等等。
比如,具体对应到上面GPU数据并行训练的场景,节点 表示每个GPU,数组 表示每个GPU各自计算的梯度,数组 表示梯度的聚合结果:
定义完问题之后,我们来讨论下这个问题的算法设计。一个常见的算法思路,是所有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.