Efficient Sparse Collective Communication and its application to Accelerate Distributed Deep Learnin

试图解决的问题

首先介绍一下什么是 Collective Communication：在 MPI 编程中，一般使用三种集体交流方式（Broadcast, AllReduce, AllGather），他们往往是将多个处理器之间的信息进行交换。

为了提高处理器间的通信速度，进而加速分布式深度学习，本篇文章提出了 OmniReduce 方法。其利用了信息的稀疏性，只将非零块进行传输。

背景介绍

最普遍的分布式深度学习（DDL）算法是随机梯度下降（SGD）。分布式 SGD 分为两个步骤：(1) 每个 worker 训练一个模型的局部拷贝，并行地处理训练数据的不同子集。(2) 将所用的 worker 的计算结果组合，得到一个平均的梯度，并将模型参数进行更新。

在 DDL 中，通信负载占比很高，严重影响了其并行化的训练速度，是一个众所周知的性能瓶颈。

目前，深度学习模型的梯度稀疏性越来越高，其中 DeepLight，LSTM 两个深度神经网络的梯度稀疏度超过了 94%。

方法介绍

OmniReduce 把输入数据分成块，这种块要么是一个输入向量连续值的分割，要么是一个表示非零值的 key-value 对的列表。

OmniReduce由两个组件 worker 和 aggregator 组成。aggregator 指导 workers 下一步要发送哪些块，为了扩展性，aggregator 可能会有多个节点，每个节点有多个槽，每个槽可聚集一块数据。worker 负责检测和发送非零块。

初始算法

对于只有一个节点一个槽的 aggregator 的 OmniReduce，有如下算法：

对于 worker，启动时首先发送第一个块，然后开始等待，直到 aggregator 给自己发送了聚合后的块，将聚合后的块接收之后，判断一下 aggregator 需要的下一个块的块号（包含在接受的数据中）和自己想要发的下一块的块号（下一个非零块）是否相同，如果相同，就发送，不同则不发。注意，worker 发送的数据中有下一个非零块的块号，aggregator 根据此块号决定下一次需要哪些块。如果 worker 没有下一步要发的数据了，就把想要发的下一块的块号置为正无穷。

对于 aggregator，每次收到 worker 发来的数据，都要把块中的值加到槽中，并把 worker 下次要发的块号记录在 next 数组（每个 worker 对应一个索引）中。如果当前收到的块号小于 next 数组的最小值，就意味着其他 worker 也已经发送过此块（或者其他 worker 的此块是零值），也就是说当前块已经收集完毕，可以将聚合后的数据发给所有的 worker。注意，此时发送的数据中有 aggregator 下次需要的块号（next 数组的最小值），供 worker 进行比较判断是否需要发送数据块。如果 aggregator 收到的 worker 想要发的下一块的块号均为正无穷，就说明可以结束聚合操作了。

论文算法截图：

细粒度并行化

很容易观察到，上面描述的聚合逻辑能够跨槽并行化。每个槽都是一个独立的聚合单元。实际上，可用的网络带宽限制了槽的数量，这些槽可在 aggregator 响应之前被并行寻址。

OmniReduce 利用了这种细粒度的并行化实现了一种流水线的方式来提高性能。aggregator 维护了一个可通过索引寻址的槽池。workers 为独立的聚合流运行算法1，每个聚合流使用一个单独的槽。当包被网络序列化，整个架构能够被视为一种基于流水线的处理方式。可用的网络带宽规定了流水线的深度。

Block Fusion

在之前的基本解决方案中，每个包包含一个块，整个的负载就是块所占的字节数加上元数据。更大的块增加了带宽的利用率，但是也导致了更小的块稀疏度。

为了平衡块稀疏度和带宽利用率，本文提出了一种叫做 Block Fusion 的方法，这个方法将多个块打包成一个包，使用一个槽成批量的处理它们。和基本的解决方案一样，一个槽保持着最小的聚合单元，有能力聚合一个包里的所有值。因此，如果 w 是融合到一个包里的块的数量，那么一个槽一次能聚合 w * b 个值。

主要的技术挑战是如何选择哪些块去融合到一个包里。虽然让每个 worker 聚合它之后的 w 个非零块看起来是很直接的，但是注意到当 workers 同时发送相同偏移的块时，流聚合时最有效率的。因此，我们必须确保每个块都要基于块偏移被映射到包中的一致的位置。因此，一个 worker 中的梯度张量被排列为一个块的二维矩阵。
比如下图这个例子：

然后，我们基于块的列索引把每个在特定位置的块映射到包中，然后我们在每个包中包含了后面的非零块索引。对于每个独立的列，通过扫描该列所在的行来找到下一个块。这种方案确保了同一个列的两个块不能被融合到同一个包中。因此，这种方案保持了和基本解决方案同样的逻辑，不需要传输其他的信息。

当 aggregator 收到了一个包后，它使用和基本解决方案同样的方式来检查每个 worker 的块是否都已经接受完毕（算法1第22行）。在一个槽中的所有块都已经聚合完毕后，aggregator 把聚合后的数据多播给所有的 worker。worker 一收到结果包，就检查下一个被 aggregator 需要的块是不是非零的。对于槽中的所有的 w 个块，aggregator 需要的下一块的块号都为正无穷时，归约就完成了。

包包含一个块数量域，指示它包含多少个块。aggregator 根据下一个块的索引值来确定每个包含的非零块的列索引。