FOAM: Blocked State Folding for Memory-Efficient LLM Training

2025-12-07

AI Summary

To address the excessive GPU memory overhead of Adam-style optimizers in large language model (LLM) training—often exceeding 70% of total memory—this paper proposes BlockAdam, a memory-efficient optimizer based on block-wise gradient mean compression and residual correction. Its core innovation is a block-state folding mechanism: optimizer states are compressed via per-parameter-block averaging, while dynamic residual correction preserves critical gradient information—avoiding computationally expensive or performance-degrading techniques such as SVD, projection, or parameter freezing. We theoretically establish that BlockAdam achieves the same convergence rate as Adam under standard non-convex assumptions. Empirical evaluation on mainstream LLM training shows that BlockAdam reduces optimizer state memory by 90%, cuts total GPU memory consumption by ~50%, and maintains or even exceeds the training throughput of full-rank Adam and existing memory-efficient baselines.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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 reduces total training memory by approximately 50%, eliminates up to 90% of optimizer state memory overhead, 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.

Problem

• Reduces memory bottlenecks in large language model training
• Compresses optimizer states to cut memory overhead significantly
• Maintains convergence rates while accelerating training speed

Innovation

• Compresses optimizer states via block-wise gradient means
• Incorporates residual correction to recover lost information
• Reduces training memory by 50% and optimizer state by 90%