[论文解读] Parallel Training of Deep Networks with Local Updates

Michael Laskin, Luke Metz+6 more|arXiv (Cornell University)|Dec 7, 2020

Advanced Neural Network Applications参考文献 66被引用 10

一句话总结

本文提出局部并行（local parallelism），一种用截断的逐层反向传播替代全局反向传播的方法，以实现深度网络的完全异步、计算高效的训练。通过在完整前向传播完成前独立更新每一层，该方法减少了通信和内存开销，在多层感知机（MLPs）上实现最高10倍的加速，在Transformer上实现2倍加速，且在ResNets上相比标准反向传播的样本吞吐量提高了40%。

ABSTRACT

Deep learning models trained on large data sets have been widely successful in both vision and language domains. As state-of-the-art deep learning architectures have continued to grow in parameter count so have the compute budgets and times required to train them, increasing the need for compute-efficient methods that parallelize training. Two common approaches to parallelize the training of deep networks have been data and model parallelism. While useful, data and model parallelism suffer from diminishing returns in terms of compute efficiency for large batch sizes. In this paper, we investigate how to continue scaling compute efficiently beyond the point of diminishing returns for large batches through local parallelism, a framework which parallelizes training of individual layers in deep networks by replacing global backpropagation with truncated layer-wise backpropagation. Local parallelism enables fully asynchronous layer-wise parallelism with a low memory footprint, and requires little communication overhead compared with model parallelism. We show results in both vision and language domains across a diverse set of architectures, and find that local parallelism is particularly effective in the high-compute regime.

研究动机与目标

• 通过探索替代性并行化策略，解决大规模小批量训练下数据并行和模型并行的收益递减问题。
• 克服反向传播中前向与反向阶段的锁定机制，该机制限制了并行性并增加了内存和同步开销。
• 探究局部逐层更新（尽管并非严格基于梯度）是否能够实现深度网络的可扩展、高效训练。
• 通过真实世界实现，在视觉和语言架构中展示局部并行的可行性和优势。
• 从通信、内存和训练效率角度，将局部并行与反向传播进行比较，量化其在高计算场景下的优势。

提出的方法

• 用截断的逐层反向传播替代全局反向传播，使每一层在部分前向传播后即可独立更新。
• 实现分块局部并行，其中每一层处理本地小批量数据，并在下游层完成前向传播前更新其权重。
• 在每一层使用辅助分类器提供局部监督信号，将权重更新与全局梯度计算解耦。
• 在层之间实施异步训练，消除反向传播锁定，支持独立的流水线计算。
• 通过仅在相邻层之间传输必要的激活和梯度数据，最小化处理器间通信，相比流水线反向传播减少约50%的总数据传输量。
• 集成激活重计算技术以减少反向传播中的内存使用，实现公平比较，同时保持局部并行的低内存占用。

实验结果

研究问题

• RQ1 与标准反向传播相比，局部逐层更新是否能在深度网络中实现显著的训练加速？
• RQ2 在通信开销和内存消耗方面，局部并行与数据并行和模型并行相比如何？
• RQ3 局部更新在多大程度上保留了与全局反向传播相似的有用特征和梯度表示？
• RQ4 在真实世界的视觉和语言模型中，特别是在高计算场景下，局部并行的实际性能提升如何？
• RQ5 局部并行是否在保持模型准确率的同时，实现了更高的硬件利用率和吞吐量？

主要发现

• 在简单MLP上，局部并行将顺序计算步骤最多减少了10倍，在Transformer架构上实现了2倍加速，证明了其在训练效率上的显著提升。
• 在ImageNet上训练ResNet34时，局部并行相比反向传播将样本吞吐量（每秒训练样本数）提高了40%，主要得益于更高的硬件利用率。
• 与流水线反向传播相比，局部并行的处理器间通信量减少了约50%，在4个IPU上总数据传输量为43.2MB，而反向传播为86.4MB。
• 在局部并行中，内存消耗始终较低或相当，尤其在高微小批量和多处理器配置下，这是由于激活存储需求减少所致。
• 尽管并非基于梯度，局部更新仍保留了与反向传播相似的特征和梯度表示，支持其在优化中的有效性。
• 激活重计算虽降低了反向传播的内存使用，但并未抵消局部并行在规模上的内存优势，尤其是在更大微小批量和更多处理器的情况下。