分布式训练 Baidu All Reduce

Baidu All Reduce，即Ring All Reduce。Ring All Reduce技术在高性能计算领域很常用，2017年被百度用于深度学习训练。

朴素All Reduce的通信时间随GPU节点数线性增长。Ring All Reduce的通信时间跟GPU节点数无关，只受限于GPU间最慢的连接。 

Ring All Reduce包含两步：scatter reduce和all gather。

1）scatter reduce：GPU交换数据，每个GPU得到最后结果的一部分（chunk）。

假设要实现数组间对应元素求和，GPU节点数为N，每个GPU都有一个相同size的数组。

1、每个GPU把自己的数组划分成N份。

2、每个GPU做N-1次迭代，每次迭代：GPU向它的右相邻节点发送一个chunk，同时从它的左相邻节点接收一个chunk并跟本地对应chunk累加。每次迭代每个GPU发送和接受的chunk都不一样。第N个GPU，一开始发送chunk N并且接收chunk N-1，然后不断向后处理。每一次迭代都把上一次迭代收到的chunk，发送出去。 

2）all gather：GPU交换chunk，每个GPU得到最后结果。

经过scatter reduce，每个GPU的某些chunk已经是最后结果了。all gather类似scatter reduce（也是N-1次迭代的发送和接收），但不是把接收到的chunk跟本地chunk累加，而是覆盖本地chunk。 第N个GPU，一开始发送第N+1个chunk，并且接收chunk N。在之后的迭代中，总是把上一次迭代接收到的chunk，发送出去。 

每个GPU的通信数据量：D = 2(N-1)S/N。

通信时间：T = 2(N-1)S/(NB)，可以看出，跟GPU节点数无关。

每次训练迭代（跟上面的迭代不是一回事），每个GPU前向传播计算误差，然后反向传播计算每个weight的梯度。反向计算梯度是从输出层的权值到输入层的权值，意味着，前面层的梯度没算出来时，后面层的梯度已经可用。因为All Reduce可以操作参数子集，所以在一些梯度还没算出来时，就可以开始All Reduce了。通信和计算之间的overlap取决于网络的实际优化情况。