MLSys’25 极低内存消耗：用SGD的内存成本实现AdamW的优化性能

作者
AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
本文第一作者 Hanqing Zhu 与 Zhengyu Zhang 均为 UT Austin 博士生。Hanqing Zhu 专注于高效 AI 计算，致力于优化机器学习硬件、系统及算法；Zhengyu Zhang 聚焦构建高效可靠的机器学习系统。通讯作者为 UT Austin 的 David Z. Pan、Zhangyang Wang 以及 Meta AI 的 Jinwon Lee 。

大语言模型（LLM）正以惊人的速度重塑各个领域的格局，但其预训练与微调过程却面临着前所未有的资源考验 —— 不仅需要庞大的计算能力，更对内存提出了极高要求。

训练 LLM 时，不仅需要存储亿级参数，还必须额外保存梯度与优化器状态（例如 Adam 中的动量和方差）。

例如，预训练 LLaMA 7B 的一个批次就消耗高达 58GB 内存：14GB 用于模型参数，42GB 用于优化器状态和梯度。这巨大的 “隐形” 内存开销迫使研究者不得不选用显存更大的高端 GPU、增加 GPU 数量，甚至牺牲训练速度调低批量大小。

首次以类 SGD 内存成本完成大模型训练

UT Austin 和 Meta AI 推出了全新训练策略 ——APOLLO（Approximated Gradient Scaling for Memory Efficient LLM Optimization）。这一创新方法首次实现了以 SGD 级别的内存成本训练大模型，同时保持了 AdamW 的卓越优化性能（已在 LLaMA 7B 的预训练上验证）。

文章已经被 MLSys 2025 接受，同时目前 APOLLO 已在 Hugging Face Transformers、LLaMA-Factory 等主流开源项目中落地，研究者可以轻松调用这一方法开展低内存大模型训练与微调。

核心亮点

  • 极低内存消耗：首次以类 SGD 内存成本完成大模型训练，达到甚至超越 AdamW 的性能。

  • 无需 SVD 计算：首次实现仅需轻量级随机投影进行大模型预训练，甚至在 7B 模型上优化速度超越 Adam。

卓越系统性能

  • 3 倍预训练加速：在 8 块 A100 GPU 上，APOLLO 预训练 LLaMA 7B 模型实现了 3 倍的加速。

  • 突破规模限制：首次利用 DDP 成功训练 13B 模型，并在 12GB 内存的消费级 GPU（如 NVIDIA RTX 4090）上完成 7B 模型的预训练，无需依赖模型并行、检查点或卸载策略。

  • 论文地址：https://arxiv.org/pdf/2412.05270
  • 论文网站：https://zhuhanqing.github.io/APOLLO/
  • 论文代码: https://github.com/zhuhanqing/APOLLO

方法简要介绍

从 Adam 到结构化更新

在传统的 Adam 优化器中，学习率的自适应更新是基于每个参数的梯度动量和方差：


APOLLO 首先将 Adam 的更新过程等效理解为对原始梯度  施加一个基于梯度动量和方差的梯度缩放因子，进一步发现这一缩放因子可以通过结构化更新（例如按通道或张量进行缩放）来近似实现。

在 LLaMA-130M 上的实验验证表明，结构化更新不仅不会导致性能损失，反而会略微提升模型性能（见图 1）。这一发现为 APOLLO 的设计提供了理论基础。

图 1: 结构化梯度缩放与 Adam 的训练损失对比

APOLLO 框架

尽管结构化更新简化了优化器设计，但仍需保存完整的优化器状态，内存开销依然较高。这促使解决以下关键问题：

问题①：是否可以将结构化的学习率自适应转换为一种既实用又内存高效的优化方法？

作者提出了 APOLLO，在低秩空间中构造一个辅助优化器状态，用于预测结构化更新的梯度缩放因子（见图 2）。由于优化器状态建立在低秩空间中，其内存占用显著减少。

图 2: APOLLO 框架

问题②：低秩空间中的自适应学习率能否有效近似其在原始空间中的行为？

APOLLO 通过理论证明，低秩空间中的结构化梯度缩放因子能够有效近似其在原始空间中的行为。同时从低秩空间预测的缩放因子会以 √(n/r) 的比例变化（n 为原始空间维度，r 为低秩空间的秩）。

这使得 APOLLO 可通过放缩缩放因子来弥补低秩带来的误差， 从而可采用极低的秩近似，在不牺牲性能的前提下实现极低的内存消耗。实验表明，在 LLaMA-7B 上，APOLLO 仅需 256 的秩，性能仍优于使用 1024 秩的 GaLore。

图 3：梯度缩放因子随秩的变化按照理论推导的 √(n/r) 变化

问题③：是否仍需要昂贵的 SVD 来构造低秩空间？

不需要。APOLLO 无需传统方法中基于昂贵 SVD 运算的低秩投影方法，只需采用基于随机投影的方法。这种方法避免了 SVD 的计算开销，引入了极低的额外计算，甚至 APOLLO 在 7B 上的优化速度还快于 Adam。

问题④：是否能进一步压缩优化器状态至 SGD 的内存消耗，同时达到甚至超越 AdamW 的性能？

为实现极致内存高效，APOLLO 提出了 APOLLO-Mini。APOLLO-Mini 通过仅使用 rank-1 的辅助子空间进行张量级梯度缩放，而非 APOLLO 中的通道级更新。

实验结果显示，APOLLO-Mini 在实现 SGD 内存消耗的同时，其预训练性能甚至超越了 AdamW。这一结果充分证明了结构化学习率更新的有效性，并为资源有限的研究者提供了高效训练大模型的解决方案。

下图为将 APOLLO 应用于 Adam 的算法：


实验及结果

研究者对 APOLLO 的 LLM 预训练、微调以及系统级性能优势进行了全面评估。

预训练

研究团队在 C4 数据集（经过净化处理的 Common Crawl 语料库）上对基于 LLaMA 的大型语言模型进行了预训练实验。

实验结果显示 APOLLO 系列方法（即使在使用极低秩的情况下，如秩分别为 256 和 1）均能显著降低内存需求，并且获得比 Adam 以及当前 SOTA 方法更好的结果。

表 1: 使用 C4 数据集预训练不同大小的 LLaMA 模型验证集困惑度（↓）对比

在 LLaMA 7B 模型上，APOLLO 系列不仅在内存使用上取得巨大突破，其训练性能也远超过 Adam 以及使用高秩（例如 1024）的 GaLore 方法。

表 2: 使用 C4 数据集预训练 LLaMA-7B 模型验证集困惑度（↓）对比

微调

微调任务（常识推理和 MMLU）中，APOLLO 的表现同样抢眼，达到了甚至超越 Adam 基线的效果。


系统优势

  • 加速训练
在硬件资源受限的情况下，APOLLO 能够支持比 AdamW 高达 4 倍的批量大小，从而实现最高可达 3 倍的训练加速。

  • 极低内存消耗

首次实现了在 12GB 内存的消费级 GPU 上成功预训练 LLaMA 7B 模型，为大模型训练在低资源环境下提供了新的可能性。

图 4：APOLLO 内存和训练速度优势

  • 极低的计算开销

由于不依赖昂贵的 SVD 计算，APOLLO 优化时间在 7B 上甚至比 adam 还要快。
随着大语言模型规模不断攀升，内存需求也在呈爆炸式增长。APOLLO 不仅在理论上打破了优化器内存瓶颈，更在实践中展现出预训练与微调的卓越性能。APOLLO 为资源受限的团队提供了一条低成本、高效率的大模型训练之路，预示着未来大模型训练将进入一个全新的时代。

© 

（文：机器之心）

发表评论