用极简优化器实现低内存大模型预训练，性能媲美甚至超越主流方法。
Memory-Efficient LLM Pretraining via Minimalist Optimizer Design
- 仅对输出层梯度做列归一化和一阶动量，大幅降低内存开销。
- 在60M-1B模型上达到与Adam相当的性能，内存仅需35%-45%。
- 适合资源受限场景的大规模语言模型训练，尤其适合7B级模型。
大语言模型训练依赖自适应优化器（如Adam），其需额外存储一阶、二阶动量，导致显著内存开销。尽管已有如GaLore、Fira、APOLLO等压缩状态的内存高效变体，但一个根本问题仍待解答：对普通SGD进行最小化修改，能否达到顶尖预训练性能？本文采用自下而上的方法系统探究此问题，提出两项简单却高效的技巧：(1) 沿输出维度进行梯度列归一化，可在无动量情况下提升SGD性能；(2) 仅在输出层应用一阶动量，因该层梯度方差最大。二者结合形成SCALE（Stochastic Column-normalized Last-layer Momentum）优化器，用于内存高效预训练。在多个模型（60M-1B）上，SCALE性能匹配或超越Adam，仅使用35%-45%总内存。且持续优于GaLore、Fira、APOLLO等现有内存高效优化器。在LLaMA 7B上，其困惑度更低、内存消耗更少，优于当前最优方法APOLLO和Muon。
原文摘要 · Abstract (English)
Training large language models (LLMs) relies on adaptive optimizers such as Adam, which introduce extra operations and require significantly more memory to maintain first- and second-order moments than SGD. While recent works such as GaLore, Fira and APOLLO have proposed state-compressed memory-efficient variants, a fundamental question remains: What are the minimum modifications to plain SGD needed to match state-of-the-art pretraining performance? We systematically investigate this question using a bottom-up approach, and identify two simple yet highly (memory- and compute-) efficient techniques: (1) column-wise gradient normalization (normalizing the gradient along the output dimension), that boosts SGD performance without momentum; and (2) applying first-order momentum only to the output layer, where gradient variance is highest. Combining these two techniques lead to SCALE (Stochastic Column-normAlized Last-layer momEntum), a simple optimizer for memory efficient pretraining. Across multiple models (60M-1B), SCALE matches or exceeds the performance of Adam while using only 35-45% of the total memory. It also consistently outperforms memory-efficient optimizers such as GaLore, Fira and APOLLO, making it a strong candidate for large-scale pretraining under memory constraints. For LLaMA 7B, SCALE outperforms the state-of-the-art memory-efficient methods APOLLO and Muon in both perplexity and memory consumption.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。