基于草图的通信高效分布式SGD

神经网络的大规模分布式训练常常受限于网络带宽，其中通信时间远超本地计算时间。受草图方法在次线性/流式算法中成功的启发，我们提出Sketched SGD，一种通过通信草图而非完整梯度来执行分布式SGD的算法。我们证明Sketched SGD在几类函数上具有优良的收敛速率。当考虑所有通信——包括梯度和更新后的模型权重——时，Sketched SGD将所需通信量从其他梯度压缩方法的$\mathcal{O}(d)$或$\mathcal{O}(W)$降低到$\mathcal{O}(\log d)$，其中$d$为模型参数数量，$W$为参与训练的worker数量。我们在transformer模型、LSTM和残差网络上运行实验，展示了总通信成本最高降低40倍且最终模型性能无损失。我们还通过实验表明Sketched SGD可扩展至至少256个worker而不增加通信成本或降低模型性能。