[论文解读] GACT: Activation Compressed Training for Generic Network Architectures

Xiaoxuan Liu, Lianmin Zheng+10 more|arXiv (Cornell University)|Jun 22, 2022

Advanced Neural Network Applications被引用 5

一句话总结

GACT 是一种通用的激活压缩训练（ACT）框架，通过动态压缩激活并提供收敛性保证，实现了对各类神经网络架构——卷积网络、Transformer 和图网络——的内存高效训练。它将激活内存减少高达 8.1 倍，使在相同 GPU 内存预算下批量大小最高可扩大 24.7 倍，且精度损失可忽略不计。该框架已作为 PyTorch 库实现，可无缝集成至现有模型及内存节省技术中。

ABSTRACT

Training large neural network (NN) models requires extensive memory resources, and Activation Compressed Training (ACT) is a promising approach to reduce training memory footprint. This paper presents GACT, an ACT framework to support a broad range of machine learning tasks for generic NN architectures with limited domain knowledge. By analyzing a linearized version of ACT's approximate gradient, we prove the convergence of GACT without prior knowledge on operator type or model architecture. To make training stable, we propose an algorithm that decides the compression ratio for each tensor by estimating its impact on the gradient at run time. We implement GACT as a PyTorch library that readily applies to any NN architecture. GACT reduces the activation memory for convolutional NNs, transformers, and graph NNs by up to 8.1x, enabling training with a 4.2x to 24.7x larger batch size, with negligible accuracy loss. We implement GACT as a PyTorch library at https://github.com/LiuXiaoxuanPKU/GACT-ICML.

研究动机与目标

• 开发一种通用的激活压缩框架，用于在不依赖架构或算子特异性假设的前提下训练各类神经网络架构。
• 证明在不预先了解算子类型或模型结构的情况下，激活压缩训练仍具有收敛性。
• 设计一种自适应算法，在运行时根据梯度影响估计结果，为每个张量选择最优压缩比。
• 将 GACT 实现为 PyTorch 库，使其可无缝集成至现有训练流程，并兼容其他内存效率技术（如梯度检查点和交换技术）。

提出的方法

• 将 ACT 建模为算子计算图，以实现跨架构的泛化，避免对层的特定假设。
• 使用线性化近似分析 ACT 的随机梯度，证明其为无偏估计，并具有结构化方差，从而支持数值估计。
• 开发一种数值算法，用于预测不同压缩策略下的梯度方差，从而实现最优策略选择。
• 利用整数规划近似方法，基于张量对梯度更新的敏感度，自动确定每个张量的最优压缩比。
• 实现 GACT 为 PyTorch 库，支持多级优化，兼容用户自定义算子，并可与现有内存效率技术集成。
• 通过在训练过程中估计每个张量对梯度方差的影响，实现在运行时动态调整压缩比。

实验结果

研究问题

• RQ1 是否可以在不依赖架构假设的前提下，使激活压缩训练在通用神经网络架构上实现可证明收敛？
• RQ2 如何自适应地为每个张量选择压缩比，以维持训练的稳定性和精度？
• RQ3 能否构建一个通用的 ACT 框架，支持任意算子（包括用户自定义算子），而无需进行定制化实现？
• RQ4 在通用深度学习模型中，内存减少、训练速度与精度损失之间的权衡关系如何？
• RQ5 GACT 如何与现有内存效率技术（如梯度检查点和交换）集成并增强其效果？

主要发现

• GACT 在卷积网络、Transformer 和图神经网络上均实现高达 8.1 倍的激活内存减少，且精度损失可忽略不计。
• 在相同 GPU 内存预算下，GACT 支持将批量大小最高扩大 24.7 倍，显著提升训练吞吐量。
• 在 ResNet-152 上，GACT 允许在 16GB GPU 内存下实现模型深度提升 7.0 倍、宽度扩大 3.6 倍或分辨率提高 3.0 倍。
• 对于 BERT-large，GACT 在相同内存约束下支持模型深度提升 2.0 倍或宽度扩大 1.6 倍。
• 将 GACT 与梯度检查点结合，在批量大小为 288 时，峰值激活内存减少 5.4 倍，且无精度下降。
• GACT 与交换技术结合后，训练速度最高提升 2.3 倍，原因在于压缩张量在 CPU 与 GPU 间通信更快；预取机制额外带来约 7% 的速度提升，且内存开销极低。