Memory-Efficient LLM Pretraining via Minimalist Optimizer Design

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. Code is available at https://github.com/OptimAI-Lab/MinimalistLLMPretraining.

Training modern language models requires huge amounts of computer memory. A large part of this memory comes from the extra information used by existing training methods, such as the Adam optimizer. This makes it expensive and difficult to train larger models on smaller computing systems. Instead of starting from a complex optimizer and trying to reduce its memory usage, we begin with the simplest method, which updates the model using the training signal (gradient) directly (stochastic gradient descent), and try to add only the minimum components needed. We found that two simple and lightweight ideas worked surprisingly well: normalizing the training gradients separately for each column of the model, and smoothing the gradient only in the last layer, where it fluctuates the most. Using the above ideas on top of stochastic gradient descent, our method, called SCALE, matched or outperformed widely used training methods while using far less memory. It also performed better than several recent memory-saving approaches, including GaLore, Fira, and APOLLO. Our results suggest that large language models can be trained efficiently without depending on complicated training algorithms.