arXiv:2412.05270v4 通知类型: replace-cross 摘要：大规模语言模型（LLMs）在训练过程中通常非常耗内存，尤其是在使用流行的AdamW优化器时。这种内存负担要求使用更多或更高性能的GPU或者减少批量大小，这限制了训练的可扩展性和吞吐量。为了应对这一问题，已经提出了各种节省内存的优化器来减少优化器的内存使用。然而，它们面临关键挑战：（i）依赖昂贵的SVD操作；（ii）与AdamW相比，性能存在显著折衷；（iii）仍然有大量优化器内存开销以维持竞争力。 在这项工作中，我们发现可以将AdamW的学习率调整规则有效地简化为结构化学习率更新。基于这一认识，我们提出了Approximated Gradient Scaling for Memory-Efficient LLM Optimization（APOLLO），它使用基于纯随机投影的辅助低秩优化器状态来近似学习率缩放。这种结构化学习率更新规则使APOLLO能够容忍进一步的内存减少，同时提供可比的预训练性能。即使其秩1变体APOLLO-Mini，在与SGD水平内存成本下，其预训练性能也优于AdamW。 广泛的经验研究表明，APOLLO系列与AdamW表现持平或更好，同时通过几乎消除AdamW的优化状态而实现更大的内存节省。这些节省带来了系统级优势：（1）增强吞吐量：在8xA100-80GB设置下，与AdamW相比，吞吐量提高3倍，支持4倍更大的批量大小。 （2）提高模型可扩展性：在A100-80GB GPU上使用原始DDP预训练LLaMA-13B，无需系统级优化。 （3）低性能GPU友好预训练：使用不到12 GB的内存进行权重量化，在单个GPU上预训练LLaMA-7B。