大型语言模型 (LLMs) 正在改变我们的世界，但其高昂的训练成本一直是阻碍其发展的重要因素。为了解决这一问题，微软研究人员提出了一种名为 FP8-LM 的新框架，利用 FP8 低精度训练，将训练成本降低了 50%。
FP8：下一代低精度训练的突破
FP8 是下一代低精度数据类型，相比传统的 FP16 和 FP32，它可以实现两倍的速度提升，同时将内存和通信成本降低 50% 至 75%。然而，目前对 FP8 训练的支持非常有限，仅有的框架是 Nvidia Transformer Engine (TE)，但它只将 FP8 应用于矩阵乘法计算，其他操作仍然使用高精度，因此无法充分发挥 FP8 的优势。
FP8-LM：全方位 FP8 训练框架
FP8-LM 框架的核心思想是将 FP8 计算、存储和通信渗透到整个 LLM 训练过程中，实现前向和后向传递都使用 FP8，从而大幅减少系统工作负载。
具体而言，FP8-LM 提供了三个优化级别，逐步整合 FP8 到混合精度和分布式训练中：
- O1 优化： 将 FP8 应用于权重和权重梯度，并在通信中使用 FP8。
- O2 优化： 支持 Adam 和 AdamW 两种优化器使用 FP8。
- O3 优化： 支持 FP8 用于分布式并行训练和 ZeRO 优化器，这对于训练大型模型至关重要。
挑战与解决方案
使用 FP8 训练 LLM 面临着数据下溢或上溢等挑战，以及由于 FP8 数据格式的动态范围和精度有限而产生的量化误差。为了解决这些问题，FP8-LM 提出了两种技术：
- 精度解耦： 将数据精度对权重、梯度、优化器状态等参数的影响解耦，并将较低的精度分配给对精度不敏感的组件。
- 自动缩放： 通过动态调整张量缩放因子，保持梯度值在 FP8 数据格式的表示范围内，从而减轻下溢和上溢情况。
显著的性能提升
实验结果表明，FP8-LM 在各种模型规模上都取得了显著的性能提升：
- 内存使用量减少： 与 BF16 混合精度训练相比，FP8-LM 将内存使用量减少了 27% 至 42%，例如 GPT-7B 减少了 27%，GPT-175B 减少了 42%。
- 通信开销降低： FP8-LM 将权重梯度通信开销降低了 63% 至 65%。
- 训练时间缩短： 在 GPT-175B 模型训练中，FP8-LM 将训练时间减少了 17%，同时内存消耗减少了 21%。
广泛的适用性
FP8-LM 框架不仅适用于预训练阶段，也适用于下游微调任务，包括指令调整和人类偏好对齐。
- 指令调整： FP8-LM 在指令调整任务中取得了与 BF16 相当的性能，同时训练速度提高了 27%。
- 强化学习与人类反馈 (RLHF)： FP8-LM 在 RLHF 中也展现出巨大潜力，可以减少模型权重和优化器状态的内存消耗。
未来展望
FP8-LM 的发布为下一代低精度训练系统树立了新的典范，它将推动大型基础模型的进一步发展。未来，研究人员将继续探索 FP8-LM 的应用，例如：
- 训练更大的 FP8 GPT 模型。
- 将 FP8-LM 应用于多模态大型模型训练。
- 在边缘设备上部署 FP8 LLM。
参考文献
总结： FP8-LM 框架是 LLM 训练领域的一项重大突破，它为降低训练成本、提高训练效率开辟了新的道路，将推动下一代大型基础模型的快速发展。