FOAM: Blocked State Folding for Memory-Efficient LLM Training

TL;DR

FOAM is a memory-efficient optimizer for training large language models that significantly reduces memory usage while maintaining convergence and performance.

Contribution

The paper introduces FOAM, a novel optimizer that compresses optimizer states with block-wise gradient means and residual correction, achieving memory savings without performance loss.

Findings

Eliminates up to 90% of optimizer memory overhead.

Accelerates convergence compared to standard Adam.

Compatible with other memory-efficient optimizers, matching or surpassing their performance.

Abstract

Large language models (LLMs) have demonstrated remarkable performance due to their large parameter counts and extensive training data. However, their scale leads to significant memory bottlenecks during training, especially when using memory-intensive optimizers like Adam. Existing memory-efficient approaches often rely on techniques such as singular value decomposition (SVD), projections, or weight freezing, which can introduce substantial computational overhead, require additional memory for projections, or degrade model performance. In this paper, we propose Folded Optimizer with Approximate Moment (FOAM), a method that compresses optimizer states by computing block-wise gradient means and incorporates a residual correction to recover lost information. Theoretically, FOAM achieves convergence rates equivalent to vanilla Adam under standard non-convex optimization settings. Empirically, FOAM eliminates up to 90\% of the memory overhead of optimizer states and accelerates convergence. Furthermore, FOAM is compatible with other memory-efficient optimizers, delivering performance and throughput that match or surpass both full-rank and existing memory-efficient baselines. Code is available at https://github.com/zqOuO/FOAM.