COAT：用于内存高效FP8训练的压缩优化器状态与激活
机器学习
2025-02-14 v3 人工智能
摘要
FP8训练已成为提高训练效率的有前景方法。现有框架通过对线性层应用FP8计算来加速训练，同时将优化器状态和激活保留在更高精度，这未能充分优化内存使用。本文引入COAT (压缩优化器状态和激活用于FP8训练)，一种新的FP8训练框架，旨在显著减少训练大型模型时的内存占用。COAT通过两个关键创新措施解决了当前局限性：(1) 动态范围扩展，将优化器状态分布与FP8表示范围更 closely 对齐，从而减小量化误差；(2) 混合粒度激活量化，通过结合per-tensor和per-group量化策略来优化激活内存。实验表明，COAT在各种任务中均能有效地将端到端训练内存占用降低约1.54倍，同时实现近乎无损的性能，包括大型语言模型预训练和微调以及视觉语言模型训练。COAT还实现了相对于BF16的1.43倍端到端训练加速，性能与或超过TransformerEngine的加速相当。COAT使得在更少的GPU上高效训练大型模型成为可能，并在分布式训练设置中实现批量大小的翻倍，为大规模模型训练提供了实用解决方案。代码已公开于 https://github.com/NVlabs/COAT。
引用
@article{arxiv.2410.19313,
title = {COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training},
author = {Haocheng Xi and Han Cai and Ligeng Zhu and Yao Lu and Kurt Keutzer and Jianfei Chen and Song Han},
journal= {arXiv preprint arXiv:2410.19313},
year = {2025}
}
备注
Accepted by ICLR 2025. 22 pages. 9 Figures. 13 Tables