一种面向低带宽网络的全局 Top-$k$ 稀疏化分布式同步 SGD 算法

分布式同步随机梯度下降（S-SGD）已广泛用于训练大规模深度神经网络（DNN），但它通常需要在计算节点（如 GPU）间以极高通信带宽迭代交换梯度。近来，Top-$k$ 稀疏化技术被提出用于减少节点间需交换的数据量。Top-$k$ 稀疏化可在不影响模型收敛的前提下将大部分梯度置零。然而，稀疏梯度须连同其不规则索引一并传输，这使得稀疏梯度的聚合十分困难。当前使用 AllGather 累加稀疏梯度的方法通信复杂度为 $O(kP)$（$P$ 为节点数），在节点数众多、带宽较低的网络上效率低下。我们观察到模型更新并不需要来自 $P$ 个节点的所有 top-$k$ 梯度，因而提出一种新颖的全局 Top-$k$（gTop-$k$）稀疏化机制来解决该问题。具体而言，我们从 $P$ 个节点中选取全局绝对值最大的 top-$k$ 梯度，而非每轮迭代累加所有局部 top-$k$ 梯度来更新模型。基于 gTop-$k$ 稀疏化的梯度聚合方法将通信复杂度从 $O(kP)$ 降至 $O(k\log P)$。在不同 DNN 上的大量实验验证了 gTop-$k$ S-SGD 的收敛表现与 S-SGD 几乎一致，且泛化性能仅有轻微下降。在扩展效率方面，我们在一个由 32 台 GPU 机器通过 1 Gbps 以太网互联的集群上评估 gTop-$k$。实验结果表明，相较 S-SGD，我们的方法取得了 $2.7-12\times$ 更高的扩展效率，相较现有 Top-$k$ S-SGD 有 $1.1-1.7\times$ 的提升。