用随机投影近似学习率更新,实现类似SGD的内存开销、媲美AdamW的训练效果。
APOLLO: SGD-like Memory, AdamW-level Performance
- 以结构化学习率更新替代传统AdamW,用低秩随机投影降低状态存储。
- rank-1版本在仅需SGD级内存下,预训练性能优于AdamW,支持4倍大批次。
- 适合资源受限场景,可在单卡低显存上训练70亿参数模型。
大语言模型训练对内存要求极高,尤其使用AdamW优化器时。现有内存高效优化器面临三大挑战:依赖昂贵的SVD计算、性能显著低于AdamW,且仍需较大优化器状态内存。本文发现,AdamW的学习率自适应规则可被粗粒度化为结构化更新。基于此,提出APOLLO:通过纯随机投影构建辅助低秩状态,近似学习率缩放。该结构化更新使APOLLO对内存压缩极度鲁棒,同时保持与AdamW相当的预训练性能。其最小版本APOLLO-Mini仅需SGD级内存,性能却超越AdamW。实验证明,APOLLO系列在性能上持平或优于AdamW,同时几乎消除AdamW的优化器状态开销。系统级收益显著:(1)在8xA100-80GB配置下,吞吐提升3倍,支持4倍大批次;(2)无需系统级优化即可在原生DDP下预训练LLaMA-13B;(3)结合量化可在单卡<12GB显存下训练LLaMA-7B。
原文摘要 · Abstract (English)
Large language models (LLMs) are notoriously memory-intensive during training, particularly with the popular AdamW optimizer. This memory burden necessitates using more or higher-end GPUs or reducing batch sizes, limiting training scalability and throughput. To address this, various memory-efficient optimizers have been proposed to reduce optimizer memory usage. However, they face critical challenges: (i) reliance on costly SVD operations; (ii) significant performance trade-offs compared to AdamW; and (iii) still substantial optimizer memory overhead to maintain competitive performance. In this work, we identify that AdamW's learning rate adaptation rule can be effectively coarsened as a structured learning rate update. Based on this insight, we propose Approximated Gradient Scaling for Memory-Efficient LLM Optimization (APOLLO), which approximates learning rate scaling using an auxiliary low-rank optimizer state based on pure random projection. This structured learning rate update rule makes APOLLO highly tolerant to further memory reductions while delivering comparable pre-training performance. Even its rank-1 variant, APOLLO-Mini, achieves superior pre-training performance compared to AdamW with SGD-level memory costs. Extensive experiments demonstrate that the APOLLO series performs on-par with or better than AdamW, while achieving greater memory savings by nearly eliminating the optimization states of AdamW. These savings provide significant system-level benefits: (1) Enhanced Throughput: 3x throughput on an 8xA100-80GB setup compared to AdamW by supporting 4x larger batch sizes. (2) Improved Model Scalability: Pre-training LLaMA-13B with naive DDP on A100-80GB GPUs without system-level optimizations. (3) Low-End GPU Friendly Pre-training: Pre-training LLaMA-7B on a single GPU using less than 12 GB of memory with weight quantization.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。