新型FP8训练框架COAT：显著减少大模型训练中的内存占用，同时保持高性能

共计 1113 个字符，预计需要花费 3 分钟才能阅读完成。

加州大学伯克利分校、英伟达、麻省理工学院和清华大学的研究人员共同提出了一种名为 COAT 的新型 FP8 训练框架。COAT 旨在显著减少大模型训练中的内存占用，同时保持高性能。这一框架通过两项关键创新解决了现有方法的局限性，从而在多种任务中实现了显著的改进。

• 论文：https://arxiv.org/abs/2410.19313
• GitHub：https://github.com/NVlabs/COAT

例如，我们正在训练一个具有数十亿参数的大语言模型（LLM），如自然语言处理领域的 GPT 或 BERT 模型。在传统的训练方法中，这些模型需要大量的内存来存储优化器状态（如梯度和动量）和中间激活值。COAT 框架通过将这些数据压缩到 FP8 格式，显著减少了内存占用，使得在有限的硬件资源下训练更大的模型成为可能。

现有框架的局限性

现有的 FP8 训练框架主要通过将 FP8 计算应用于线性层来加速训练，但优化器状态和激活仍然保持在较高的精度（如 FP32 或 BF16）。这种做法未能充分利用内存使用，限制了训练效率的提升。

COAT 的关键创新

1. 动态范围扩展：
   • 目的：使优化器状态分布更接近 FP8 表示范围，从而减少量化误差。
   • 方法：通过动态调整优化器状态的范围，使其更适合 FP8 表示。这一技术有效地减少了量化过程中的信息损失，提高了模型的训练精度。
2. 混合粒度激活量化：
   • 目的：优化激活内存的使用，减少内存占用。
   • 方法：结合张量级和组级量化策略，灵活地选择最适合当前任务的量化粒度。这种混合策略不仅减少了内存占用，还保持了激活的精度。

实验结果

实验表明，COAT 在多个方面表现优异：

• 内存占用：相比 BF16，COAT 有效地减少了 1.54 倍的端到端训练内存占用。
• 性能：在各种任务中，如大型语言模型的预训练和微调、视觉语言模型的训练，COAT 实现了几乎无损的性能。
• 训练加速：COAT 实现了相比 BF16 的 1.43 倍端到端训练加速，表现与 TransformerEngine 的加速相当或超越。

应用前景

COAT 的推出为大规模模型训练提供了几个重要的优势：

• 高效的全参数训练：COAT 使得在更少的 GPU 上进行高效的全参数训练成为可能，降低了硬件成本。
• 批量大小翻倍：在分布式训练设置中，COAT 能够将批量大小翻倍，进一步提高训练效率。
• 扩展性：COAT 为扩展大规模模型训练提供了一个实用的解决方案，适用于多种深度学习任务和应用场景。

总结

COAT 通过动态范围扩展和混合粒度激活量化两项关键技术，显著提高了 FP8 训练的效率和性能。这一框架不仅减少了内存占用，还保持了高性能，为大规模模型训练提供了一个实用且高效的解决方案。未来，随着这一技术的进一步发展和应用，我们有理由期待更多创新的深度学习模型和应用。