[论文解读] PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization

Thijs Vogels, Sai Praneeth Karimireddy+1 more|arXiv (Cornell University)|May 31, 2019

Advanced Data Compression Techniques参考文献 35被引用 96

一句话总结

PowerSGD 引入基于幂迭代、带误差反馈与 all-reduce 聚合的低秩梯度压缩器，在分布式训练中实现接近 SGD 的准确性，同时带来显著的通信加速。

ABSTRACT

We study gradient compression methods to alleviate the communication bottleneck in data-parallel distributed optimization. Despite the significant attention received, current compression schemes either do not scale well or fail to achieve the target test accuracy. We propose a new low-rank gradient compressor based on power iteration that can i) compress gradients rapidly, ii) efficiently aggregate the compressed gradients using all-reduce, and iii) achieve test performance on par with SGD. The proposed algorithm is the only method evaluated that achieves consistent wall-clock speedups when benchmarked against regular SGD with an optimized communication backend. We demonstrate reduced training times for convolutional networks as well as LSTMs on common datasets. Our code is available at https://github.com/epfml/powersgd.

研究动机与目标

• 动机与解决深度学习数据并行分布式优化中的通信瓶颈。
• 开发可扩展的线性梯度压缩器，支持 all-reduce 聚合。
• 通过误差反馈实现有偏压缩的收敛性并保持测试精度。
• 在不同数据集和硬件上，展示相对于全精度 SGD 的墙钟时间加速，适用于 CNNs 与 LSTMs。
• 提供一个实用、可扩展的方法及开源实现。

提出的方法

• 引入一个秩-r 梯度压缩器，将每个梯度矩阵 M 近似为 PQ^T，使用从上一步热启动的一步子空间/幂迭代。
• 利用线性性实现跨 W 个工作节点的 all-reduce 聚合。
• 应用误差反馈并在后压缩时引入动量以实现有偏压缩的收敛。
• 将每一层的梯度分解成矩阵并独立压缩每个矩阵，令 P ∈ R^{n×r}，Q ∈ R^{m×r，r 较小（1–4）}。
• 将 PQ^T 进行解压缩，并使用带动量的分布式更新（EF-SGD with Momentum）进行。
• 在 CIFAR-10/ResNet18 和 Wikitext-2/LSTM 上经验性地与 SGD 及其他压缩器（Signum、Spectral Atomo）进行对比，并报告墙钟时间、传输数据量和准确率。

实验结果

研究问题

• RQ1 基于幂迭代的低秩梯度压缩器是否能在不牺牲测试准确率的前提下实现显著的通信量下降？
• RQ2 误差反馈是否能在非凸深度学习设置中实现有偏压缩方案的收敛性和鲁棒性？
• RQ3 对幂迭代进行热启动如何影响近似质量和训练准确率随时间的变化？
• RQ4 在 CNN 和 RNN/LSTM 模型中，秩、通信开销与端到端训练时间之间的权衡是什么？
• RQ5 PowerSGD 如何随工作节点数量增加以及在不同通信后端（NCCL 与 GLOO）下扩展？

主要发现

算法	测试准确率	每轮传输的数据	每批次时间
SGD	94.3%	1023 MB	312 ms
Rank-1 PowerSGD	93.6%	4 MB	229 ms
Rank-2 PowerSGD	94.4%	8 MB	239 ms
Rank-4 PowerSGD	94.5%	14 MB	260 ms

• PowerSGD 在 16-GPU 设置下使用优化的 NCCL 后端实现相对于常规 SGD 的墙钟时间加速，同时将梯度压缩超过 120×。
• 在 CIFAR-10 的 ResNet-18 上，Rank-2 PowerSGD 在每轮通信 8 MB 的情况下达到 94.4% 的测试准确率，而 SGD 为 94.3%，通信量为 1023 MB。
• 在 Wikitext-2 的 LSTM 语言模型中，Rank-4 PowerSGD 达到与 SGD 相近的困惑度，同时将每轮发送的数据量减少约 99%，端到端训练时间减少约 55%。
• 带误差反馈的 PowerSGD 在测试准确率上优于无偏的基于秩的压缩（例如，在相似设置下，偏置 Rank-2 达到 94.4%，无偏约 75–76%）。
• 对幂迭代进行热启动可以缩小到最佳秩-r 近似的差距，在避免 SVD 代价的同时维持或提高测试准确性。
• 该方法通过 all-reduce 随着工作节点数量增加而具备良好扩展性，在优化后的后端上在速度和可扩展性方面优于替代方案（Signum）。