Research map / Large language models / Instruction and fine-tuning
BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes BAdam, a block coordinate descent method that updates one transformer block at a time with Adam, enabling memory-efficient full-parameter finetuning of Llama 3-8B on a single RTX3090 and Llama 3-70B on 4xA100.
Problem
Full-parameter finetuning of LLMs with Adam requires roughly 18M GB of GPU memory, which is prohibitive under limited GPU memory; PEFT methods like LoRA restrict updates to a low-rank subspace.
Approach
Combine block coordinate descent with Adam's update rule: partition parameters into D blocks (e.g., transformer layers), run K Adam steps on one active block with the others frozen, clear optimizer states on switching, and keep optimizer states only for the active block, giving memory around 2M + 16M/D GB.
Key results
- On Llama 3-8B MT-bench (lr 1e-6), BAdam averages 6.65 vs LoRA 6.27 and Adam 6.23.
- Memory for finetuning Llama 3-8B on a single RTX3090: BAdam 23.5GB vs LOMO 21.5GB and LoRA-rank8 22.3GB, while Adam is estimated at 144.8GB+.
- On MathInstruct, Llama 3-70B with BAdam averages 62.7 vs LoRA 59.0 across six math benchmarks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use BAdam instead of LoRA | Llama 3-8B math finetuning · Average math score | score: 43.3 -> 44.4 | Helped |
| Use BAdam instead of Adam | Llama 3-8B math finetuning · Average math score | score: 44.1 -> 44.4 | Helped |
| Use BAdam instead of LoRA | Llama 3-70B math finetuning · Average math score | score: 59.0 -> 62.7 | Helped |
| Use BAdam instead of Adam | Llama 2-7B instruction tuning · MT-bench average, learning rate 1e-5 | MT-bench score: 4.76 -> 4.96 | Helped |
| Use BAdam instead of LoRA | Llama 3-8B instruction tuning · MT-bench average, learning rate 1e-6 | MT-bench score: 6.27 -> 6.65 | Helped |
| Use BAdam instead of Adam | Llama 3-8B instruction tuning · MT-bench average, learning rate 1e-6 | MT-bench score: 6.23 -> 6.65 | Helped |
| Use BAdam instead of LoRA | Llama 3-8B finetuning · Alpaca-GPT4, per epoch | backward time: 3.20 hours -> 1.74 hours | Helped |
| Replace Adam's update rule with SGD in the BCD framework (BSGD) (ablation) | BCD finetuning on Llama 3-8B · MT-bench | MT-bench score | No clear effect |
| Use BCD updates rather than full-model updates (ablation) | SGD finetuning on Llama 3-8B · MT-bench | MT-bench score | Helped |
| Use BCD updates rather than full-model updates (ablation) | Adam finetuning on Llama 3-8B · Training loss convergence | training loss | Hurt |
| Change block ordering among ascending, descending, and random reshuffling (ablation) | BAdam on Llama 3-8B · Alpaca-GPT4 | training loss convergence | No clear effect |
| Vary the inner Adam steps K among 10, 50, 100, and 200 (ablation) | BAdam on Llama 3-8B · Alpaca-GPT4 | training loss convergence | No clear effect |
Surprising findings
- BSGD (block SGD) substantially outperforms full SGD on MT-bench despite similar convergence, suggesting block updates better preserve pretrained knowledge.
- BAdam outperforms Adam on Llama 3-8B MT-bench, although Adam at lr 1e-5 overfits Alpaca-GPT4 and was omitted.
Limitations
- Only supervised finetuning evaluated; preference optimization not explored.
- Continued pretraining experiment is preliminary.
- Convergence theory covers only the deterministic gradient case.
Experimental setup
Datasets: Alpaca-GPT4 MathInstruct SuperGLUE StarCoder-Python
Benchmarks: MT-bench GSM8K Aqua MMLU-Math SAT-Math MATH NumGLUE SuperGLUE
Metrics: MT-bench score zero-shot accuracy peak memory (GB) wall-clock time cumulative explained variance / effective rank
Baselines: Adam LoRA LOMO GaLore BSGD SGD
Builds on: Adam LoRA LOMO GaLore gradient checkpointing
Closest papers by meaning
- FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training (ICML 2025 · similarity 0.89)
- Search for Efficient Large Language Models (NeurIPS 2024 · similarity 0.89)
- Gradient Weight-normalized Low-rank Projection for Efficient LLM Training (AAAI 2025 · similarity 0.88)
- GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection (ICML 2024 · similarity 0.88)
- Second-Order Fine-Tuning without Pain for LLMs: A Hessian Informed Zeroth-Order Optimizer (ICLR 2025 · similarity 0.88)
- Puzzle: Distillation-Based NAS for Inference-Optimized LLMs (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.