一种基于误差有界有损压缩的新型内存高效深度学习训练框架

由于对预测精度和分析质量的需求不断增长，深度神经网络正变得日益加深、加宽且非线性化。在训练 DNN 模型时，中间激活数据必须在前向传播期间保存在内存中，并在反向传播时恢复。然而，由于硬件设计的限制，GPU 等最先进的加速器仅配备非常有限的内存容量，这在训练大规模 DNN 时显著限制了最大批大小，进而限制了性能加速。在本文中，我们提出了一种新型内存驱动的高性能 DNN 训练框架，该框架利用误差有界有损压缩显著减少训练的内存需求，从而允许训练更大的网络。与采用 JPEG 等基于图像的有损压缩器来压缩激活数据的最先进解决方案不同，我们的框架专门设计了具有严格误差控制机制的误差有界有损压缩。具体而言，我们提供了从改变的激活数据到梯度的压缩误差传播的理论分析，然后实证研究了改变的梯度对整个训练过程的影响。基于这些分析，我们提出了一种改进的有损压缩器和一种自适应方案，以动态配置有损压缩误差界限并调整训练批大小，从而进一步利用节省的内存空间实现额外加速。我们使用四种流行的 DNN 和 ImageNet 数据集对设计进行了评估，并与最先进的解决方案进行了对比。结果表明，与基线训练和带有压缩的最先进框架相比，我们提出的框架可以显著减少训练内存消耗，分别高达 13.5 倍和 1.8 倍，且几乎没有或完全没有精度损失。