arXiv:2410.19313v3 宣布类型: replace-cross 摘要: FP8训练已成为提高训练效率的一种有前途的方法。现有的框架通过将FP8计算应用于线性层，同时保留优化器状态和激活的高精度，从而加速训练，但这未能充分利用内存使用。本文介绍了COAT（压缩优化器状态和激活的FP8训练框架），这是一种新型的FP8训练框架，旨在在训练大规模模型时显著减少内存占用。COAT通过两项关键创新解决了当前的限制：(1) 动态范围扩展，这使得优化器状态分布更紧密地与FP8表示范围对齐，从而减少量化误差；(2) 混合粒度激活量化，这通过结合张量级和分组量化策略优化激活内存。实验表明，与BF16相比，COAT可以将端到端的训练内存占用减少1.54倍，同时在各种任务（如大型语言模型预训练和微调以及视觉语言模型训练）中几乎实现无损性能。COAT还实现了与BF16相比1.43倍的端到端训练加速，与TransformerEngine的加速相当或超过。COAT使在较少的GPU上高效训练大规模模型成为可能，并在分布式训练设置中使批次大小翻倍成为可能，提供了一种扩展大规模模型训练的实用方案。相关代码可在https://github.com/NVlabs/COAT获得。