COAT: Compressing Optimizer states and Activation for Memory-Efficient
  FP8 Training
FP8 training has emerged as a promising method for improving training
efficiency. Existing frameworks accelerate training by applying FP8 computation
to linear layers while leaving optimizer states and activations in higher
precision, which fails to fully optimize memory usage. This paper introduces
COAT (Compressing Optimizer States and Activations for FP8 Training), a novel
FP8 training framework designed to significantly reduce memory footprint when
training large models. COAT addresses current limitations through two key
innovations: (1) Dynamic Range Expansion, which aligns optimizer state
distributions more closely with the FP8 representation range, thereby reducing
quantization error, and (2) Mixed-Granularity Activation Quantization, which
optimizes activation memory using a combination of per-tensor and per-group
quantization strategies. Experiments demonstrate that COAT effectively reduces
end-to-end training memory footprint by 1.54x compared to BF16 while achieving
nearly lossless performance across various tasks, such as Large Language Model
pretraining and fine-tuning and Vision Language Model training. COAT also
achieves a 1.43x end-to-end training speedup compared to BF16, performing on
par with or surpassing TransformerEngine's speedup. COAT enables efficient
full-parameter training of large models on fewer GPUs, and facilitates doubling
the batch size in distributed training settings, providing a practical solution
for scaling large-scale model training. The code is available at
https://github.com/NVlabs/COAT.
67206ecd5ff1582170e13f7b
Haocheng Xi
claimed_verified
2025-05-27T07:57:45.136Z
66ce751a8ec9fda2cf5a9e85
/avatars/c17093ca81dad007b3e50bae503955a7.svg
Haocheng Xi
user
xihc-ucb
67206ecd5ff1582170e13f7c
Han Cai
claimed_verified
2024-11-04T21:24:19.002Z
650e2b14c945dfc9386a7e28
/avatars/0a9be20aa53f4c52a2d1a8b02d4093ea.svg
Han Cai
user
han-cai
67206ecd5ff1582170e13f7d
Ligeng Zhu
claimed_verified
2024-12-30T19:38:00.316Z
62b4b5beb25cb80fcf278354
https://cdn-avatars.huggingface.co/v1/production/uploads/62b4b5beb25cb80fcf278354/7uGNOeE91paZikhK2QFgp.jpeg
Ligeng Zhu
user
Ligeng-Zhu
67206ecd5ff1582170e13f7e
Yao Lu
67206ecd5ff1582170e13f7f
Kurt Keutzer
67206ecd5ff1582170e13f80
Jianfei Chen
67206ecd5ff1582170e13f81
Song Han
COAT: Compressing Optimizer states and Activation for Memory-Efficient
  FP8 Training
FP8 training has emerged as a promising method for improving training
efficiency. Existing frameworks accelerate training by applying FP8 computation
to linear layers while leaving optimizer states and activations in higher
precision, which fails to fully optimize memory usage. This paper introduces
COAT (Compressing Optimizer States and Activations for FP8 Training), a novel
FP8 training framework designed to significantly reduce memory footprint when
training large models. COAT addresses current limitations through two key
innovations: (1) Dynamic Range Expansion, which aligns optimizer state
distributions more closely with the FP8 representation range, thereby reducing
quantization error, and (2) Mixed-Granularity Activation Quantization, which
optimizes activation memory using a combination of per-tensor and per-group
quantization strategies. Experiments demonstrate that COAT effectively reduces
end-to-end training memory footprint by 1.54x compared to BF16 while achieving
nearly lossless performance across various tasks, such as Large Language Model
pretraining and fine-tuning and Vision Language Model training. COAT also
achieves a 1.43x end-to-end training speedup compared to BF16, performing on
par with or surpassing TransformerEngine's speedup. COAT enables efficient
full-parameter training of large models on fewer GPUs, and facilitates doubling
the batch size in distributed training settings, providing a practical solution
for scaling large-scale model training. The code is available at
https://github.com/NVlabs/COAT.
COAT reduces memory usage and accelerates training of large models through dynamic range expansion and mixed-granularity activation quantization in FP8 training.
FP8 training COAT Dynamic Range Expansion Mixed-Granularity Activation Quantization BF16 Large Language Model Vision Language Model TransformerEngine