模型并行训练中的激活与梯度压缩

大型神经网络需要庞大的机器计算集群。模型并行训练，即模型架构在多个工作节点间顺序划分，是训练现代模型的一种流行方法。信息压缩可用于减少工作节点间的通信时间，因为通信往往是此类系统的瓶颈。本文探讨了在模型并行分布式训练设置中同时压缩激活和梯度对收敛性的影响。我们分析了量化与 TopK 压缩等压缩方法，并实验了误差补偿技术。此外，我们采用了结合 AQ-SGD 逐批误差反馈方法的 TopK 压缩。我们在图像分类和语言模型微调任务上进行了实验。研究结果表明，梯度所需的压缩程度比激活更温和。我们观察到 $K=10\%$ 是最低的 TopK 压缩水平，不会严重损害模型收敛。实验还表明，使用 TopK 训练的模型仅在推理阶段也应用压缩时表现良好。我们发现，与普通压缩相比，误差反馈技术并未改善模型并行训练，但允许模型在无压缩的情况下进行推理，且质量几乎不下降。最后，当与 AQ-SGD 方法结合使用时，强于 $K=30\%$ 的 TopK 压缩会显著恶化模型性能。