Reproduce “FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training” on OpenResearch

About this paper

Philip Zmushko, Aleksandr Beznosikov, Martin Takáč, Samuel Horváth

FRUGAL introduces a memory-efficient optimization framework for training large language models by splitting gradients and using state-full optimizers for a small subspace and state-free optimizers for the rest. This approach enables full-rank updates with significantly reduced optimizer state memory, achieving performance comparable to full AdamW training.

Key insights

• Full-rank updates are crucial for effectively leveraging all gradient information and enabling comprehensive exploration of the parameter space, leading to improved model performance, especially during pre-training.
• The output layer and classification head in Transformer models are highly sensitive to the choice of optimizer, performing best with advanced, state-full methods like AdamW, while other modules can be effectively optimized with state-free methods.
• Using signSGD as the state-free optimizer for the majority of parameters in the complementary subspace offers strong performance and memory efficiency, often outperforming plain SGD due to its update characteristics.

Summary by alphaXiv.