[论文解读] Communication-efficient distributed SGD with Sketching

Nikita Ivkin, Daniel Rothchild+4 more|arXiv (Cornell University)|Mar 12, 2019

Stochastic Gradient Optimization Techniques参考文献 37被引用 16

一句话总结

该论文提出 Sketched-SGD，一种通信高效的分布式随机梯度下降算法，通过传输大小为 𝒪(log d) 的梯度草图而非完整梯度，实现次线性通信复杂度。在 Transformer、LSTM 和残差网络上，该方法在不损失模型准确率的前提下，将总通信成本降低高达 40 倍，并且在不降低性能的前提下可有效扩展至 256 个工作节点。

ABSTRACT

Large-scale distributed training of neural networks is often limited by network bandwidth, wherein the communication time overwhelms the local computation time. Motivated by the success of sketching methods in sub-linear/streaming algorithms, we introduce Sketched SGD, an algorithm for carrying out distributed SGD by communicating sketches instead of full gradients. We show that Sketched SGD has favorable convergence rates on several classes of functions. When considering all communication -- both of gradients and of updated model weights -- Sketched SGD reduces the amount of communication required compared to other gradient compression methods from $\mathcal{O}(d)$ or $\mathcal{O}(W)$ to $\mathcal{O}(\log d)$, where $d$ is the number of model parameters and $W$ is the number of workers participating in training. We run experiments on a transformer model, an LSTM, and a residual network, demonstrating up to a 40x reduction in total communication cost with no loss in final model performance. We also show experimentally that Sketched SGD scales to at least 256 workers without increasing communication cost or degrading model performance.

研究动机与目标

• 解决大规模分布式深度学习中的通信瓶颈问题，其中网络带宽限制了训练速度。
• 将每个工作节点的通信成本从 𝒪(d) 和 𝒪(W) 降低至 𝒪(log d)，从而实现在大量工作节点上的高效扩展。
• 开发一种方法，在实现高压缩率的同时保持收敛性保证和模型性能。
• 克服现有梯度压缩技术在工作节点数量（W）增加时扩展性差的局限性。
• 利用流式算法中的草图技术，实现通信开销极低的实用且可扩展的分布式训练。

提出的方法

• 使用 Count Sketch 和稀疏恢复技术，将梯度压缩为每个工作节点大小为 𝒪(log d) 的草图。
• 仅传输本地梯度的草图，而非完整梯度，将通信量从 𝒪(d) 降低至 𝒪(log d)。
• 在参数服务器端使用草图恢复算法，以有界误差重建近似完整的梯度。
• 通过确保草图保留足够信息以支持有效优化，维持收敛性特性。
• 将草图机制集成到标准的同步数据并行 SGD 中，保持标准训练工作流。
• 支持无偏和有偏的梯度估计，并在适当假设下提供理论收敛保证。

实验结果

研究问题

• RQ1 草图技术能否在不损失收敛性的前提下，将分布式 SGD 的通信成本降低至 𝒪(d) 以下？
• RQ2 所提出的方法在不增加通信成本或性能损失的前提下，能否在大量工作节点（W）下实现高效扩展？
• RQ3 草图技术能否在真实世界模型上实现高压缩比（例如 40 倍），且不降低最终模型的准确率？
• RQ4 与现有梯度压缩方法（如 top-k、量化）相比，Sketched-SGD 在通信效率和可扩展性方面表现如何？
• RQ5 在草图引起的梯度近似下，Sketched-SGD 的理论收敛行为如何？

主要发现

• Sketched-SGD 在 Transformer、LSTM 和残差网络模型上将总通信成本最高降低 40 倍，且最终模型性能无损失。
• 该方法在不增加通信成本或降低模型准确率的前提下，可有效扩展至至少 256 个工作节点。
• 每个工作节点的通信成本降低至 𝒪(log d)，在模型大小 d 上为次线性，在工作节点数量 W 上为常数。
• MNIST 实验表明，Sketched-SGD 的训练和测试误差率与原始 SGD 一致，符合理论收敛速率。
• 在小规模实验中，草图大小为 280（40 列，7 行）时，k=10 且 P=10 的条件下压缩比约为 4。
• 理论分析表明，在光滑凸函数的标准假设下，Sketched-SGD 保持了有利的收敛速率。