arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2305.19982cs.LGcs.AI
收藏

用于大规模深度神经网络训练的 Adam 累积以减少激活和梯度内存占用

Adam Accumulation to Reduce Memory Footprints of both Activations and Gradients for Large-scale DNN Training

Shanghai Jiao Tong University(上海交通大学);Microsoft Research Asia(微软亚洲研究院)

Yijia Zhang, Yibo Han, Shijie Cao, Guohao Dai, Youshan Miao, Ting Cao, Fan Yang, Ningyi Xu

AI总结:

本文提出 Adam Accumulation(AdamA)方法,通过将梯度直接整合进优化器状态并在微批次上累积,同时减少激活与梯度内存,在保持 Adam 收敛性的前提下实现高达 23% 的内存降低,并支持训练更大规模模型。

摘要:

GPU 内存耗尽已成为大规模深度神经网络训练的主要瓶颈。如何在训练过程中减少内存占用受到了广泛的研究关注。我们发现,以往的梯度累积方法虽然减少了激活内存,但由于保留梯度与释放梯度之间的矛盾,无法与梯度内存的减少相兼容。为了解决这一问题,我们提出了一种针对 Adam 的新型优化器累积方法,命名为 Adam Accumulation(AdamA),它能够同时减少激活内存和梯度内存。具体而言,AdamA 直接将梯度整合到优化器状态中,并在微批次上累积优化器状态,从而使梯度在使用后可以立即释放。我们从数学上和实验上证明了 AdamA 具有与 Adam 相同的收敛性质。在基于 Transformer 的模型上评估时,与梯度累积相比,AdamA 实现了高达 23% 的内存减少,同时训练吞吐量的下降不到 2%。值得注意的是,AdamA 可以与针对优化器状态的内存减少方法协同工作,在不同内存容量的 GPU 上,相较于 PyTorch 和 DeepSpeed 基线,能够容纳大 1.26 倍至 3.14 倍的模型。

英文摘要:

Running out of GPU memory has become a main bottleneck for large-scale DNN training. How to reduce the memory footprint during training has received intensive research attention. We find that previous gradient accumulation reduces activation memory but fails to be compatible with gradient memory reduction due to a contradiction between preserving gradients and releasing gradients. To address this issue, we propose a novel optimizer accumulation method for Adam, named Adam Accumulation (AdamA), which enables reducing both activation and gradient memory. Specifically, AdamA directly integrates gradients into optimizer states and accumulates optimizer states over micro-batches, so that gradients can be released immediately after use. We mathematically and experimentally demonstrate AdamA yields the same convergence properties as Adam. Evaluated on transformer-based models, AdamA achieves up to 23% memory reduction compared to gradient accumulation with less than 2% degradation in training throughput. Notably, AdamA can work together with memory reduction methods for optimizer states to fit 1.26x~3.14x larger models over PyTorch and DeepSpeed baseline on GPUs with different memory capacities.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑