Research map / Architectures and efficiency / Quantization
COAT: Compressing Optimizer states and Activations for Memory-Efficient FP8 Training
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.
COAT quantizes AdamW optimizer states and activations to FP8 (via dynamic range expansion and mixed-granularity activation quantization) to cut training memory by 1.54x and speed up training 1.43x versus BF16 with near-lossless accuracy.
Problem
FP8 training frameworks such as Transformer Engine and FP8-LM leave activations and the second-order optimizer momentum in higher precision, so memory savings are incomplete, and naive FP8 quantization of these tensors loses accuracy.
Approach
Apply a power-function expansion f(x)=sign(x)|x|^k per quantization group so optimizer-state dynamic ranges fill the E4M3 range, and quantize activations with fine-grained (1x16) groups for non-linear layers and per-tensor scaling (via two-stage Group Scaling) for linear layers, storing activations directly in FP8.
Key results
- End-to-end memory reduction of 1.54x vs BF16 and up to 1.43x end-to-end speedup on Llama models, with OLMo-1B average accuracy 67.8% vs BF16 68.2%.
- Dynamic Range Expansion reduces quantization MSE of m/(sqrt(v)+eps) by 1.63x for the optimizer states.
- Single transformer layer activation memory reduced 1.65x vs BF16, with up to 1.57x speedup; COAT enables full-parameter Llama-2-7B training on a single GPU where BF16 and TE OOM.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply COAT instead of BF16 training | OLMo-1B pretraining · OLMo-1B pretraining, Avg Acc | accuracy: 68.2 % -> 67.8 % | Hurt |
| Apply COAT instead of BF16 training | OLMo-7B pretraining · OLMo-7B pretraining, Avg ppl | perplexity: 11.174 -> 11.279 | Hurt |
| Apply COAT instead of BF16 fine-tuning | Llama-2-7B math fine-tuning · Math fine-tuning | average score: 55.7 -> 55.5 | No clear effect |
| Use COAT rather than BF16 | Llama-2-7B end-to-end training on 4 GPUs · Context length 2048 | peak memory reduction ratio: 1.00× -> 1.54× | Helped |
| Quantize optimizer states and activations to FP8 rather than use BF16 | Llama-2-7B training on 4 GPUs · End-to-end training | speedup: 1.00x -> 1.45x | Helped |
| Add Dynamic Range Expansion to optimizer-state quantization (ablation) | FP8 optimizer states · Optimizer quantization | MSE: +1.63× MSE reduction | Helped |
| Add Dynamic Range Expansion to DE8 optimizer quantization (ablation) | DE8 optimizer states · Optimizer quantization | quantization error: 1.41× | Helped |
| Quantize activations to FP8 with COAT's activation quantization flow (ablation) | LLM training · OLMo-1B downstream tasks | average accuracy: 51.6 -> 52.3 | No clear effect |
| Quantize optimizer states to FP8 (ablation) | Llama-2-7B on 4 GPUs · End-to-end training | speedup: 1.00x -> 1.08x | Helped |
| Add FP8 activation quantization to FP8 optimizer quantization (ablation) | Llama-2-7B on 4 GPUs · End-to-end training | speedup: 1.08x -> 1.45x | Helped |
| Use per-group quantization rather than per-tensor quantization for linear layers (ablation) | Linear-layer activations · Linear layers | speed: 1300 TFlops -> about 900 TFlops | Hurt |
Surprising findings
- Per-group quantization of linear layers reduces linear-layer speed by about 40%, so the authors use per-tensor scaling for linear layers despite its lower accuracy granularity.
- The paper reports the main accuracy degradation comes from optimizer-state quantization, while activation quantization causes almost no degradation.
Limitations
- OLMo-7B COAT results show notable drops on some tasks (e.g., ARC-Easy 61.9 vs 65.7 BF16), and average accuracy drops slightly (72.7 vs 73.2).
- Fine-tuning and VLM results show small drops on some benchmarks (e.g., GSM8K 56.6 vs 57.7), and the evaluation is limited to a few model sizes.
- Dynamic range expansion requires computing k on the fly per group per step, adding complexity; FSDP statistics are not synchronized across GPUs.
Experimental setup
Datasets: Dolma MAmmoTH VILA SFT data mixture
Benchmarks: WikiText-103 C4 Pile COPA ARC-Easy SciQ HellaSwag MathQA SVAMP NumGLUE GSM8K VideoMME
Metrics: perplexity accuracy memory footprint (GB) tokens/s throughput MSE of quantization error
Baselines: BF16 TransformerEngine FP8-LM
Builds on: Transformer Engine FP8-LM 8-bit Adam Jetfire VS-Quant
Closest papers by meaning
- Scaling FP8 training to trillion-token LLMs (ICLR 2025 · similarity 0.87)
- Towards Fully FP8 GEMM LLM Training at Scale (NeurIPS 2025 · similarity 0.86)
- FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training (ICML 2025 · similarity 0.86)
- µnit Scaling: Simple and Scalable FP8 LLM Training (ICML 2025 · similarity 0.85)
- MGUP: A Momentum-Gradient Alignment Update Policy for Stochastic Optimization (NeurIPS 2025 · similarity 0.85)
- Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation (ICML 2024 · similarity 0.85)
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.