FP8-LM:训练成本降低50%的LLM新时代

大型语言模型 (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 训练领域的一项重大突破,它为降低训练成本、提高训练效率开辟了新的道路,将推动下一代大型基础模型的快速发展。

发表评论