本文提出 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.
The Chinese University of Hong Kong, Shenzhen(香港中文大学(深圳));Shenzhen International Center for Industrial and Applied Mathematics(深圳国际工业与应用数学中心);Shenzhen Research Institute of Big Data(深圳大数据研究院);Duke University(杜克大学);Anthropic(Anthropic公司);Stanford University(斯坦福大学)
AdamS: Momentum Itself Can Be A Normalizer for LLM Pretraining and Post-training
Wangxuan Institute of Computer Technology, Peking University(北京大学王宣计算机技术研究所);University of Science and Technology of China(中国科学技术大学);State Key Laboratory of General Artificial Intelligence(通用人工智能国家重点实验室)
A Theoretical and Empirical Study on the Convergence of Adam with an "Exact" Constant Step Size in Non-Convex Settings
Robert Bosch Centre for Cyber-Physical Systems(罗伯特·博世网络物理系统中心);Indian Institute of Science(印度科学学院);Netaji Subhas University of Technology(纳塔吉·苏巴斯·技术大学);TCS Research(塔塔咨询服务公司研究中心)
本文研究非凸优化中 Adam 的固定步长收敛问题,推导并实现“精确”恒定步长,给出确定性与随机版本的收敛保证及运行时间界,并通过分类实验验证其降低梯度范数的效果。
Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching
Mila & DIRO, Université de Montréal(蒙特利尔大学米拉与DIRO);Critical ML, University of Waterloo(滑铁卢大学Critical ML);Zhejiang University(浙江大学);Canada CIFAR AI Chair(加拿大CIFAR AI主席)
GCond: Gradient Conflict Resolution via Accumulation-based Stabilization for Large-Scale Multi-Task Learning
Moscow Institute of Physics and Technology(莫斯科物理技术学院);Skolkovo Institute of Science and Technology(斯科尔科沃科学技术学院);Biomedically Informed Artificial Intelligence Laboratory (BIMAI-Lab), University of Sharjah(沙迦大学生物医学人工智能实验室(BIMAI-Lab))
A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models
Beijing International Center for Mathematical Research, Peking University(北京大学北京数学科学中心);Center for Data Science, Peking University(北京大学数据科学中心);Center for Machine Learning Research, Peking University(北京大学机器学习研究中心)
Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation
University of Science and Technology of China(中国科学技术大学);NLPR & MAIS, Institute of Automation, Chinese Academy of Sciences(中国科学院自动化研究所);School of Artificial Intelligence, University of Chinese Academy of Sciences(中国科学院大学人工智能学院)