浅谈 Tensorflow 分布式架构：ring all-reduce 算法

首先还是先回顾下参数同步更新和异步更新的区别：

同步更新模式下，所有 GPU 在同一时间点与参数服务器交换、融合梯度；异步更新模式下，所有 GPU 各自独立与参数服务器通信，交换、融合梯度。

• 异步更新通信效率高速度快，但往往收敛不佳，因为一些速度慢的节点总会提供过时、错误的梯度方向。可通过上一篇介绍的 Stale Synchronous Parallel Parameter Server 方法缓解该问题。

• 同步更新通信效率低，通常训练慢，但训练收敛稳定，因为同步更新基本等同于单卡调大 的 batch size 训练。

• 但是传统的同步更新方法（各个 gpu 卡算好梯度，求和算平均的方式），在融合梯度时，会产生巨大的通信数据量，这种通信压力往往在模型参数量很大时，显得很明显。因此我们需要找到一种方法，来解决同步更新的网络瓶颈问题。其中最具代表性的一种方法就是：ring all-reduce。

##parameter server 框架下同步更新方式，网络瓶颈定量分析

这边假设有 1 个 server 端（存放参数），10 个 worker 端（计算梯度），模型是 Deep Speech 2，参数量 300M，相当于 1.2 G 的大小的内存数据（300M * sizeof(float)）。假设网络带宽 1G bytes/s （万兆网卡），10 卡同步更新，需要 10.8 s 完成参数 Send。在单 ps 节点、有限带宽环境下，通信时间随着 GPU 数量的增加而线性增长，很难想象一个 10 卡的集群每训练一个 batch 都需要等待 10 ~ 20s 来同步参数！通信时延几乎完全覆盖掉了 GPU 并行计算节节省下的计算时间。当然也可以通过一些技巧来缓解通信压力，比如增加 server 的个数。

Ring Allreduce 框架下同步更新算法

定义 GPU 集群的拓扑结构：

每个 GPU 只从左邻居接受数据、并发送数据给右邻居。

算法主要分两步：

1. scatter-reduce：会逐步交换彼此的梯度并融合，最后每个 GPU 都会包含完整融合梯度的一部分。

2. allgather：GPU 会逐步交换彼此不完整的融合梯度，最后所有 GPU 都会得到完整的融合梯度

scatter-reduce

举例：数组求和

Step1：将数组在每个 GPU 上都分块

Step2：N-1 轮的 scatter-reduce，每一轮中，每个 GPU 将自己的一个 chunk 发给右邻居，并接收左邻居发来的 chunk，并累加。

Allgather

和 scatter-reduce 操作类似，只不过将每个 chunk 里面的操作由累加值变为替换。

通信代价分析：每个 GPU 在 Scatter Reduce 阶段，接收 N-1 次数据，N 是 GPU 数量；每个 GPU 在 allgather 阶段，接收 N-1 次 数据；每个 GPU 每次发送 K/N 大小数据块，K 是总数据大小；所以，Data Transferred=2(N−1)*K/N ，随着 GPU 数量 N 增加，总传输量恒定。也就是理论上，随着 gpu 数量的增加，ring all-reduce 有线性加速能力。

下面一篇文章，将给大家介绍 tensorflow 中是如何实现 ring all-reduce 算法的。

参考文献：

https://zhuanlan.zhihu.com/p/34172340

http://andrew.gibiansky.com/

本文转载自 Alex-zhai 知乎账号。

原文链接：https://zhuanlan.zhihu.com/p/69797852