Adam 累积以减少大规模 DNN 训练中与激活和梯度相关的内存占用
机器学习
2023-06-01 v1 人工智能
摘要
GPU 内存耗尽已成为大规模 DNN 训练的主要瓶颈。如何减少训练期间的内存占用已受到大量研究关注。我们发现，先前的梯度累积减少了激活内存，但由于保留梯度与释放梯度之间的矛盾，无法与梯度内存缩减相兼容。为解决此问题，我们提出一种针对 Adam 的新型优化器累积方法，名为 Adam Accumulation（AdamA），其能够同时减少激活与梯度内存。具体而言，AdamA 直接将梯度整合进优化器状态，并在微批次上累积优化器状态，从而梯度可在使用后即刻释放。我们从数学和实验上证明 AdamA 具有与 Adam 相同的收敛性质。在基于 transformer 的模型上评估，AdamA 相较梯度累积实现了高达 23% 的内存缩减，且训练吞吐量下降小于 2%。值得注意的是，AdamA 可与针对优化器状态的内存缩减方法协同工作，在不同内存容量的 GPU 上，相较 PyTorch 和 DeepSpeed 基线可容纳大 1.26 倍至 3.14 倍的模型。
引用
@article{arxiv.2305.19982,
title = {Adam Accumulation to Reduce Memory Footprints of both Activations and Gradients for Large-scale DNN Training},
author = {Yijia Zhang and Yibo Han and Shijie Cao and Guohao Dai and Youshan Miao and Ting Cao and Fan Yang and Ningyi Xu},
journal= {arXiv preprint arXiv:2305.19982},
year = {2023}
}