[论文解读] Bandwidth Reduction using Importance Weighted Pruning on Ring AllReduce

Zehua Cheng, Zhenghua Xu|arXiv (Cornell University)|Jan 6, 2019

Advanced Neural Network Applications参考文献 19被引用 8

一句话总结

本文提出了一种基于重要性加权剪枝的环形全部归约方法，以减少大规模深度学习训练中的通信带宽。通过基于梯度幅值与参数值之比动态剪枝低重要性梯度，并使用概率索引广播，该方法在仅使用GTX 1080Ti GPU的96节点集群上，实现了AlexNet高达64倍的梯度压缩和ResNet50高达58.8倍的压缩，同时保持了模型精度。

ABSTRACT

It is inevitable to train large deep learning models on a large-scale cluster equipped with accelerators system. Deep gradient compression would highly increase the bandwidth utilization and speed up the training process but hard to implement on ring structure. In this paper, we find that redundant gradient and gradient staleness has negative effect on training. We have observed that in different epoch and different steps, the neural networks focus on updating different layers and different parameters. In order to save more communication bandwidth and preserve the accuracy on ring structure, which break the restrict as the node increase, we propose a new algorithm to measure the importance of gradients on large-scale cluster implementing ring all-reduce based on the size of the ratio of parameter calculation gradient to parameter value. Our importance weighted pruning approach achieved 64X and 58.8X of gradient compression ratio on AlexNet and ResNet50 on ImageNet. Meanwhile, in order to maintain the sparseness of the gradient propagation, we randomly broadcast the index of important gradients on each node. While the remaining nodes are ready for the index gradient and perform all-reduce update. This would speed up the convergence of the model and preserve the training accuracy.

研究动机与目标

• 为解决在环形全部归约拓扑上大规模分布式深度学习训练中日益严重的通信瓶颈问题。
• 在不牺牲模型精度的前提下减少带宽使用，尤其是在节点数量增加时。
• 克服在环形结构中深度梯度压缩的局限性，因为随着节点数增加，梯度会变得更密集。
• 仅使用消费级GPU（如GTX 1080Ti）实现大规模集群中高效、高带宽的利用。
• 通过动态、逐层的梯度重要性估计，保持收敛速度和模型性能。

提出的方法

• 使用梯度幅值与参数值之比（|∇w| / |w|）来度量梯度重要性，以识别不关键的更新。
• 基于各层间梯度重要性方差与均值之比，采用动态阈值调整策略，自适应地控制剪枝程度。
• 执行逐层重要性剪枝，仅选择性地传输高重要性梯度，其余梯度在本地累积。
• 在各节点间随机广播重要梯度的索引，以保持稀疏性并确保正确的全部归约更新。
• 对低重要性梯度采用概率更新规则，以防止延迟并维持收敛性。
• 采用预热训练和本地梯度裁剪，以在训练初期稳定优化过程。

实验结果

研究问题

• RQ1 如何在不降低模型精度的前提下，有效将梯度稀疏化应用于环形全部归约拓扑？
• RQ2 何种度量可可靠识别大规模分布式训练中的冗余或不重要梯度？
• RQ3 梯度延迟如何影响收敛性？在稀疏通信方案中如何缓解该问题？
• RQ4 尽管随着节点数增加梯度密度上升，是否仍可有意义地在环形结构中应用深度梯度压缩？
• RQ5 在动态训练场景中，何种阈值策略能最优地平衡带宽减少与模型性能？

主要发现

• 所提方法在ImageNet训练中实现了AlexNet 64倍的梯度压缩比和ResNet50 58.8倍的压缩比，显著优于固定阈值基线方法。
• 通过逐层阈值调整，该方法在ResNet50上保持了76.31%的top-1精度，略高于基线（76.09%），同时实现了47.6倍的压缩率。
• 网络I/O带宽显著降低，有效通信维持在千兆速率水平，无需Infiniband支持。
• 梯度重要性的方差与均值之比被发现是自适应阈值调整的关键指标，有助于提升收敛稳定性。
• 传统深度梯度压缩在环形结构中因梯度密度随节点数增加而失效，而所提方法通过动态剪枝克服了这一问题。
• 该方法使仅使用GTX 1080Ti GPU的96节点集群实现训练，避免了对Infiniband等高端网络设施的需求。