让大模型训练内存减半,还能提速,只用FP8就行。
COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training
- 动态扩展优化器范围+分组量化激活,精准压缩内存
- 相比BF16内存降1.54倍,训练速度提升1.43倍
- 适合想在少显存上训大模型的研究者和工程师
FP8训练已成为提升训练效率的有力方法。现有框架仅对线性层使用FP8计算,保留优化器状态和激活值的高精度,未能充分优化内存占用。本文提出COAT(Compressing Optimizer States and Activations for FP8 Training),一种新型FP8训练框架,显著降低大模型训练的内存开销。通过两项创新:(1) 动态范围扩展,使优化器状态分布更贴近FP8表示范围,减少量化误差;(2) 混合粒度激活量化,结合逐张量与逐组量化策略优化激活内存。实验表明,相较于BF16,COAT在各类任务(如大语言模型预训练、微调及视觉语言模型训练)中实现近乎无损性能的同时,端到端训练内存减少1.54倍,速度提升1.43倍,达到或超过TransformerEngine表现。该框架支持在更少GPU上进行全参数训练,并可在分布式训练中将批量大小翻倍,为大规模模型训练提供实用解决方案。代码已开源:https://github.com/NVlabs/COAT。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。