奥斯汀，德克萨斯 – 大语言模型（LLM）正以前所未有的速度重塑各个领域，然而，其预训练和微调过程对计算资源和内存提出了严峻挑战。近日，来自德克萨斯大学奥斯汀分校（UT Austin）和 Meta AI 的研究团队推出了一种名为 APOLLO（Approximated Gradient Scaling for Memory Efficient LLM Optimization）的全新训练策略，有望彻底改变这一现状。该研究成果已被机器学习系统顶级会议 MLSys 2025 接收。
LLM 训练的“内存困境”
训练 LLM 不仅需要存储数以亿计的参数，还需要额外保存梯度和优化器状态，例如 Adam 中的动量和方差。以 LLaMA 7B 为例，预训练一个批次就需要高达 58GB 的内存，其中模型参数仅占 14GB，而优化器状态和梯度则占据了剩余的 42GB。这种巨大的内存开销迫使研究人员不得不选用显存更大的高端 GPU，增加 GPU 数量，甚至牺牲训练速度来降低批量大小。
APOLLO：SGD 内存成本，AdamW 优化性能
APOLLO 的核心在于，它首次实现了以 SGD（随机梯度下降）级别的内存成本训练大模型，同时保持了 AdamW 的卓越优化性能。这意味着研究人员可以使用更少的资源，更快地训练出性能更优异的 LLM。
Hanqing Zhu 和 Zhengyu Zhang，UT Austin 的两位博士生，分别是该论文的第一作者。他们与通讯作者 David Z. Pan、Zhangyang Wang（UT Austin）以及 Jinwon Lee（Meta AI）共同致力于解决 LLM 训练中的内存瓶颈问题。
APOLLO 的三大亮点：
- 极低内存消耗： APOLLO 首次以类 SGD 内存成本完成大模型训练，达到甚至超越 AdamW 的性能。
- 无需 SVD 计算： APOLLO 首次实现仅需轻量级随机投影进行大模型预训练，甚至在 7B 模型上优化速度超越 Adam。
- 卓越系统性能：
- 3 倍预训练加速： 在 8 块 A100 GPU 上，APOLLO 预训练 LLaMA 7B 模型实现了 3 倍的加速。
- 突破规模限制： APOLLO 首次利用 DDP 成功训练 13B 模型，并在 12GB 内存的消费级 GPU（如 NVIDIA RTX 4090）上完成 7B 模型的预训练，无需依赖模型并行、检查点或卸载策略。
技术解析：结构化更新与低秩近似
APOLLO 的核心思想在于将 Adam 的更新过程等效理解为对原始梯度施加一个基于梯度动量和方差的梯度缩放因子，并进一步发现这一缩放因子可以通过结构化更新（例如按通道或张量进行缩放）来近似实现。
尽管结构化更新简化了优化器设计，但仍需保存完整的优化器状态，内存开销依然较高。为了解决这个问题，APOLLO 在低秩空间中构造一个辅助优化器状态，用于预测结构化更新的梯度缩放因子。由于优化器状态建立在低秩空间中，其内存占用显著减少。
此外，APOLLO 通过理论证明，低秩空间中的结构化梯度缩放因子能够有效近似其在原始空间中的行为。更重要的是，APOLLO 无需传统方法中基于昂贵 SVD 运算的低秩投影方法，只需采用基于随机投影的方法，避免了 SVD 的计算开销。
开源落地与未来展望
目前，APOLLO 已经在 Hugging Face Transformers、LLaMA-Factory 等主流开源项目中落地，研究人员可以轻松调用这一方法开展低内存大模型训练与微调。
APOLLO 的出现，有望显著降低 LLM 训练的门槛，加速大模型在各个领域的应用。随着技术的不断发展，我们有理由相信，未来的 LLM 训练将更加高效、经济、普及。
参考文献：
- 论文地址：https://arxiv.org/pdf/2412.05270
- 论文网站：https://zhuhanqing.github.io/APOLLO/
- 论文代码: https://github.com/zhuhanqing/APOLLO
关键词： 大语言模型，LLM，内存优化，APOLLO，机器学习，人工智能，MLSys 2025，深度学习，梯度下降，AdamW，SGD。
Views: 0