[论文解读] Pufferfish: Communication-efficient Models At No Extra Cost

Wang, Hongyi, Saurabh Agarwal+1 more|arXiv (Cornell University)|Mar 5, 2021

Advanced Neural Network Applications参考文献 89被引用 5

一句话总结

Pufferfish 是一种通信和计算高效的分布式训练框架，通过训练从全秩模型衍生的低秩、预因子化深度网络来实现通信效率。通过先用全秩网络训练一小部分时间（例如10%），然后应用SVD将其转换为低秩形式，Pufferfish 消除了压缩开销，在不损失精度的情况下，相比 PyTorch 的 DDP 实现了高达1.64倍的端到端加速，优于彩票券假设和结构化剪枝方法。

ABSTRACT

To mitigate communication overheads in distributed model training, several studies propose the use of compressed stochastic gradients, usually achieved by sparsification or quantization. Such techniques achieve high compression ratios, but in many cases incur either significant computational overheads or some accuracy loss. In this work, we present Pufferfish, a communication and computation efficient distributed training framework that incorporates the gradient compression into the model training process via training low-rank, pre-factorized deep networks. Pufferfish not only reduces communication, but also completely bypasses any computation overheads related to compression, and achieves the same accuracy as state-of-the-art, off-the-shelf deep models. Pufferfish can be directly integrated into current deep learning frameworks with minimum implementation modification. Our extensive experiments over real distributed setups, across a variety of large-scale machine learning tasks, indicate that Pufferfish achieves up to 1.64x end-to-end speedup over the latest distributed training API in PyTorch without accuracy loss. Compared to the Lottery Ticket Hypothesis models, Pufferfish leads to equally accurate, small-parameter models while avoiding the burden of "winning the lottery". Pufferfish also leads to more accurate and smaller models than SOTA structured model pruning methods.

研究动机与目标

• 解决由频繁梯度传输引起的分布式深度学习训练中的通信瓶颈。
• 消除通常与梯度压缩技术（如稀疏化和量化）相关的计算开销。
• 在不牺牲模型精度或无需复杂实现更改的情况下实现通信效率。
• 提供与现有深度学习框架（如 PyTorch）兼容的即插即用解决方案。
• 在精度和模型大小方面优于现有方法，如彩票券假设和结构化剪枝。

提出的方法

• 在总训练时间的极小部分（例如10%）内训练一个全秩深度神经网络，以稳定权重。
• 对每一层应用奇异值分解（SVD），将其分解为低秩组件，从而构建预因子化的低秩模型。
• 对生成的低秩模型继续进行微调，完成剩余的训练周期。
• 将低秩模型直接集成到训练流水线中，绕过显式的梯度压缩步骤。
• 在早期训练阶段采用全秩与低秩组件结合的混合架构，以减轻精度损失。
• 在最小修改现有分布式训练API（如 PyTorch DDP）的基础上，实现端到端效率。

实验结果

研究问题

• RQ1 能否将梯度压缩直接集成到模型架构本身，以消除计算开销，同时保持精度？
• RQ2 在部分全秩训练后通过SVD对模型进行预因子化，是否能在不造成精度下降的情况下实现通信效率？
• RQ3 该方法是否能在不依赖自定义通信原语的情况下，实现优于最先进分布式训练API的加速效果？
• RQ4 与彩票券假设相比，该方法在精度、模型大小和训练效率方面表现如何？
• RQ5 低秩因子化是否可广泛应用于各种架构和任务，而无需大量超参数调优？

主要发现

• 在真实分布式设置中，Pufferfish 相比 PyTorch 的 DistributedDataParallel (DDP) 实现了高达1.64倍的端到端加速，且无任何精度损失。
• 在 ResNet-18 与 CIFAR-10 上，Pufferfish 在速度优化设置下实现了1.16倍的每轮次加速，且计算量减少22%。
• 采用早期全秩预热的混合训练策略，相比原始低秩训练，使 VGG-19-BN (CIFAR-10) 的精度提升0.36%，ResNet-50 (ImageNet) 的精度提升4.8%。
• Pufferfish 模型的精度与彩票券假设模型相当或更优，同时避免了昂贵的彩票筛选过程。
• Pufferfish 生成的模型比最先进的结构化剪枝方法更小、更准确，展现出更高的参数效率。