APOLLO: SGD-like Memory, AdamW-level Performance

📄 arXiv: 2412.05270v4 📥 PDF

作者: Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z. Pan, Zhangyang Wang, Jinwon Lee

分类: cs.LG, cs.AI, cs.PF

发布日期: 2024-12-06 (更新: 2025-02-17)

备注: Accepted to MLSys 2025; the newest version with new experiments

💡 一句话要点

APOLLO：通过近似梯度缩放，在LLM优化中实现AdamW级性能和SGD级内存效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 优化器 内存优化 随机投影 梯度缩放

📋 核心要点

  1. 现有优化器（如AdamW）在训练LLM时内存需求高，限制了训练规模和效率，且内存优化方法常依赖SVD或牺牲性能。
  2. APOLLO通过结构化的学习率更新，利用随机投影近似梯度缩放，降低优化器状态的内存占用，同时保持甚至超越AdamW的性能。
  3. 实验表明，APOLLO在内存占用大幅降低的同时，实现了与AdamW相当甚至更好的性能，显著提升了训练吞吐量和模型可扩展性。

📝 摘要（中文）

大型语言模型（LLM）训练期间的内存需求非常高，特别是使用流行的AdamW优化器时。这种内存负担迫使人们使用更多或更高端的GPU，或者减少批量大小，从而限制了训练的可扩展性和吞吐量。为了解决这个问题，已经提出了各种内存高效的优化器来减少优化器内存的使用。然而，它们面临着严峻的挑战：（i）依赖于昂贵的SVD运算；（ii）与AdamW相比，存在显著的性能折衷；（iii）为了保持竞争性能，仍然需要大量的优化器内存开销。本文发现AdamW的学习率自适应规则可以有效地粗化为结构化的学习率更新。基于这一洞察，我们提出了用于内存高效LLM优化的近似梯度缩放（APOLLO），它使用基于纯随机投影的辅助低秩优化器状态来近似学习率缩放。这种结构化的学习率更新规则使得APOLLO能够容忍进一步的内存减少，同时提供可比的预训练性能。即使是其rank-1变体APOLLO-Mini，也能以SGD级别的内存成本实现优于AdamW的预训练性能。大量的实验表明，APOLLO系列与AdamW的性能相当或更好，同时通过几乎消除AdamW的优化状态来实现更大的内存节省。这些节省提供了显著的系统级优势：（1）增强的吞吐量：在8xA100-80GB设置上，与AdamW相比，吞吐量提高了3倍，支持4倍更大的批量大小。（2）改进的模型可扩展性：在A100-80GB GPU上使用naive DDP预训练LLaMA-13B，无需系统级优化。（3）对低端GPU友好的预训练：使用权重量化，在单个GPU上使用少于12 GB的内存预训练LLaMA-7B。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）训练过程中，由于AdamW等优化器的高内存占用，导致训练成本高昂、可扩展性受限的问题。现有内存优化方法通常依赖于计算复杂度高的SVD操作，或者在性能上与AdamW存在显著差距，无法在保证性能的同时有效降低内存消耗。

核心思路：论文的核心思路是将AdamW的学习率自适应规则进行粗化，转化为一种结构化的学习率更新方式。通过近似梯度缩放，利用低秩的辅助优化器状态来模拟AdamW的学习率调整，从而在大幅降低内存占用的同时，保持AdamW的优化性能。这种设计避免了昂贵的SVD计算，并允许进一步降低优化器状态的秩，实现更极致的内存优化。

技术框架：APOLLO的整体框架是在标准优化算法（如SGD）的基础上，引入一个低秩的辅助优化器状态。该状态通过随机投影来近似AdamW的梯度缩放过程。具体流程如下：1. 计算梯度；2. 使用随机投影将梯度投影到低维空间；3. 在低维空间更新辅助优化器状态；4. 使用辅助优化器状态近似AdamW的学习率缩放；5. 更新模型参数。

关键创新：APOLLO最重要的技术创新点在于使用随机投影来近似AdamW的学习率缩放。与直接存储和更新完整的二阶矩估计相比，随机投影可以将梯度信息压缩到低维空间，从而显著降低内存占用。此外，APOLLO的结构化学习率更新方式使其对内存缩减具有很强的鲁棒性，即使使用极低的秩（如rank-1），也能保持良好的优化性能。

关键设计：APOLLO的关键设计包括：1. 随机投影矩阵的选择：可以使用高斯随机矩阵或稀疏随机矩阵，以进一步降低计算复杂度。2. 低秩辅助优化器状态的更新规则：可以使用类似于AdamW的更新规则，但只在低维空间进行。3. 学习率缩放的近似方式：可以使用不同的近似函数来将低维状态映射到学习率缩放因子。论文中提出的APOLLO-Mini使用rank-1的辅助状态，实现了极致的内存优化。

🖼️ 关键图片

img_0

📊 实验亮点

APOLLO系列优化器在内存占用上实现了显著降低，几乎消除了AdamW的优化状态。在8xA100-80GB配置下，APOLLO实现了3倍于AdamW的吞吐量，并支持4倍更大的batch size。APOLLO成功地在A100-80GB GPU上使用naive DDP预训练了LLaMA-13B，无需额外的系统级优化。APOLLO还支持在单GPU上使用少于12GB的内存预训练LLaMA-7B，并结合了权重量化技术。

🎯 应用场景

APOLLO优化器可广泛应用于大型语言模型的预训练和微调，尤其适用于资源受限的场景，如低端GPU或边缘设备。它能降低训练成本，提高训练效率，并使更大规模的模型能够在有限的硬件资源上进行训练。此外，APOLLO还有助于推动AI技术在更多领域的应用，例如自然语言处理、计算机视觉和语音识别等。

📄 摘要（原文）

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.