Less Memory Means smaller GPUs: Backpropagation with Compressed Activations

📄 arXiv: 2409.11902v1 📥 PDF

作者: Daniel Barley, Holger Fröning

分类: cs.LG

发布日期: 2024-09-18

备注: Presented at ITEM workshop co-located with ECML PKDD 2024, Vilnius LT

💡 一句话要点

提出压缩激活的反向传播方法，降低GPU内存占用，实现更小GPU上的深度学习训练。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 反向传播 激活压缩 内存优化 深度学习 池化

📋 核心要点

  1. 现有深度学习模型训练对GPU内存需求巨大，限制了在消费级GPU上的应用，尤其是在反向传播阶段存储激活值消耗大量内存。
  2. 论文提出在反向传播过程中压缩激活图，具体方法是使用池化操作，从而减少内存占用和数据移动，而前向传播保持不变。
  3. 实验结果表明，该方法能够在保持预测精度的情况下，显著降低峰值内存消耗（降低29%），但会增加训练时间。

📝 摘要（中文）

深度神经网络（DNNs）规模的不断扩大导致计算资源需求急剧增长。许多最新的架构，尤其是大型语言模型，必须使用拥有数千个加速器（如GPU或TPU）的超级计算机进行训练。除了大量的浮点运算外，DNNs的内存占用也在爆炸式增长。相比之下，GPU架构的内存容量非常有限。即使是像某些EfficientNet变体这样相对较小的架构，也无法在合理的mini-batch size下于单个消费级GPU上进行训练。在训练期间，中间输入激活必须存储到反向传播时才能用于梯度计算。这些激活占据了大部分内存。因此，本文考虑使用池化来压缩反向传播的激活图，这可以减少内存占用和数据移动量。前向计算保持未压缩状态。我们以常见的视觉架构ResNet为例，通过实验展示了收敛性，并研究了对特征检测的影响。通过这种方法，我们能够在保持预测准确性的前提下，以更长的训练时间为代价，将峰值内存消耗降低29%。

🔬 方法详解

问题定义：深度神经网络的训练，特别是大型模型，对GPU内存提出了很高的要求。在反向传播过程中，需要存储前向传播的激活值用于计算梯度，这部分内存消耗非常显著。现有的方法要么需要更大的GPU，要么限制了batch size，影响训练效率。因此，如何在有限的GPU内存下训练更大的模型是一个关键问题。

核心思路：核心思路是在反向传播过程中对激活值进行压缩，从而减少内存占用。具体来说，论文使用池化操作来压缩激活图。由于反向传播只需要梯度信息，适当的压缩不会显著影响最终的训练效果。前向传播保持不变，以保证特征提取的质量。

技术框架：整体框架是在标准的反向传播算法中，在激活值被用于梯度计算之前，先对其进行池化压缩。具体流程如下：1. 前向传播：正常进行，保存未压缩的激活值。2. 反向传播：a. 从后往前计算梯度。b. 在需要使用激活值计算梯度时，先对相应的激活值进行池化压缩。c. 使用压缩后的激活值计算梯度。d. 继续反向传播。

关键创新：关键创新在于提出了在反向传播过程中压缩激活值的思想，并且证明了这种方法在降低内存占用的同时，可以保持模型的预测精度。与现有方法相比，该方法不需要修改网络结构，易于实现，并且可以与其他内存优化技术结合使用。

关键设计：论文中关键的设计包括：1. 池化操作的选择：论文使用了常见的池化操作，如最大池化和平均池化。2. 池化大小的选择：池化大小决定了压缩的程度，需要根据具体的模型和数据集进行调整。3. 训练策略：由于压缩激活值可能会影响训练的收敛速度，因此可能需要调整学习率等超参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在ResNet模型上，该方法能够在保持预测准确性的前提下，将峰值内存消耗降低29%。虽然训练时间有所增加，但可以通过调整学习率等超参数来缓解。该方法在不同数据集和模型上的泛化能力还有待进一步研究。

🎯 应用场景

该研究成果可应用于资源受限的场景，例如在边缘设备或消费级GPU上训练大型深度学习模型。通过降低内存需求，可以使得更多研究者和开发者能够在有限的硬件条件下进行深度学习研究和应用开发。此外，该方法还可以与其他内存优化技术结合使用，进一步降低内存占用，从而训练更大规模的模型。

📄 摘要（原文）

The ever-growing scale of deep neural networks (DNNs) has lead to an equally rapid growth in computational resource requirements. Many recent architectures, most prominently Large Language Models, have to be trained using supercomputers with thousands of accelerators, such as GPUs or TPUs. Next to the vast number of floating point operations the memory footprint of DNNs is also exploding. In contrast, GPU architectures are notoriously short on memory. Even comparatively small architectures like some EfficientNet variants cannot be trained on a single consumer-grade GPU at reasonable mini-batch sizes. During training, intermediate input activations have to be stored until backpropagation for gradient calculation. These make up the vast majority of the memory footprint. In this work we therefore consider compressing activation maps for the backward pass using pooling, which can reduce both the memory footprint and amount of data movement. The forward computation remains uncompressed. We empirically show convergence and study effects on feature detection at the example of the common vision architecture ResNet. With this approach we are able to reduce the peak memory consumption by 29% at the cost of a longer training schedule, while maintaining prediction accuracy compared to an uncompressed baseline.