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
this https URL

本文提出了 GACT，一种用于支持各种神经网络结构的 ACT 框架，旨在减小训练内存占用。通过分析 ACT 的近似梯度的线性化版本，我们证明了 GACT 的收敛性，而且不需要关于操作符类型或模型架构的先验知识。此外，我们提出了一种算法，通过估计每个张量对梯度的影响来决定压缩比率，以保证训练的稳定性。GACT 作为 PyTorch 库的实现，可用于卷积神经网络，transformers 和图神经网络，在使用 8.1 倍的较小的激活存储器减少训练内存占用，并且在几乎没有精度损失的情况下可以进行 4.2 到 24.7 倍的大批量训练。