FP8-LM: Training FP8 Large Language Models

📄 arXiv: 2310.18313v2 📥 PDF

作者: Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang, Ze Liu, Yifan Xiong, Ziyue Yang, Bolin Ni, Jingcheng Hu, Ruihang Li, Miaosen Zhang, Chen Li, Jia Ning, Ruizhe Wang, Zheng Zhang, Shuguang Liu, Joe Chau, Han Hu, Peng Cheng

分类: cs.LG, cs.CL

发布日期: 2023-10-27 (更新: 2023-12-19)

🔗 代码/项目: GITHUB

💡 一句话要点

提出FP8混合精度框架以提升大语言模型训练效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱八：物理动画 (Physics-based Animation) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: FP8 混合精度 大语言模型 训练效率 低精度训练 分布式学习 优化器状态 内存使用

📋 核心要点

  1. 现有方法在大语言模型训练中对内存和计算资源的需求较高，导致训练成本上升。
  2. 本文提出了一种FP8自动混合精度框架，能够在不影响模型精度的情况下，使用低精度数据格式进行训练。
  3. 实验结果显示，该框架在训练效率和内存使用上均有显著提升，降低了大模型的训练成本。

📝 摘要（中文）

本文探讨了FP8低位数据格式在大语言模型（LLMs）高效训练中的应用。我们的关键见解是，LLM训练中的大多数变量（如梯度和优化器状态）可以使用低精度数据格式，而不会影响模型的准确性，也无需更改超参数。我们提出了一种新的FP8自动混合精度框架，提供三种FP8利用级别，以简化LLM的混合精度和分布式并行训练。实验结果表明，在H100 GPU平台上训练GPT-175B模型时，我们的FP8混合精度训练框架不仅实现了39%的实际内存使用减少，还比广泛采用的BF16框架（即Megatron-LM）快75%，超越了Nvidia Transformer Engine的速度37%。此外，该方法具有通用性，可无缝应用于LLM指令调优和基于人类反馈的强化学习，降低微调成本。

🔬 方法详解

问题定义：本文旨在解决大语言模型训练中对内存和计算资源的高需求问题。现有的BF16框架在效率和资源使用上存在不足，导致训练成本高昂。

核心思路：论文的核心思路是通过引入FP8低精度数据格式，允许在训练过程中使用低精度表示而不损失模型的准确性。该方法通过逐步引入8位梯度和优化器状态，优化了混合精度训练的过程。

技术框架：整体架构包括三个主要模块：FP8梯度计算、优化器状态管理和分布式学习。该框架支持逐步集成FP8格式，确保训练过程的高效性和稳定性。

关键创新：最重要的技术创新点在于提出了一种新的FP8自动混合精度框架，能够在不改变超参数的情况下，显著降低内存使用和提高训练速度。这与现有的BF16方法形成了鲜明对比。

关键设计：关键设计包括对FP8格式的支持、优化器状态的低精度表示，以及在分布式训练中的高效实现。具体的参数设置和损失函数设计确保了训练过程的稳定性和高效性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，FP8混合精度训练框架在训练GPT-175B模型时，实际内存使用减少了39%，训练速度比BF16框架快75%，超越Nvidia Transformer Engine的速度37%。这些结果展示了该方法在大规模模型训练中的显著优势。

🎯 应用场景

该研究的潜在应用领域包括大语言模型的训练、指令调优以及基于人类反馈的强化学习等。通过降低训练成本和提高效率，该框架能够加速模型的开发和部署，推动AI技术的进一步应用与发展。

📄 摘要（原文）

In this paper, we explore FP8 low-bit data formats for efficient training of large language models (LLMs). Our key insight is that most variables, such as gradients and optimizer states, in LLM training can employ low-precision data formats without compromising model accuracy and requiring no changes to hyper-parameters. Specifically, we propose a new FP8 automatic mixed-precision framework for training LLMs. This framework offers three levels of FP8 utilization to streamline mixed-precision and distributed parallel training for LLMs. It gradually incorporates 8-bit gradients, optimizer states, and distributed learning in an incremental manner. Experiment results show that, during the training of GPT-175B model on H100 GPU platform, our FP8 mixed-precision training framework not only achieved a remarkable 39% reduction in real memory usage but also ran 75% faster than the widely adopted BF16 framework (i.e., Megatron-LM), surpassing the speed of Nvidia Transformer Engine by 37%. This largely reduces the training costs for large foundation models. Furthermore, our FP8 mixed-precision training methodology is generic. It can be seamlessly applied to other tasks such as LLM instruction tuning and reinforcement learning with human feedback, offering savings in fine-tuning expenses. Our FP8 low-precision training framework is open-sourced at {https://github.com/Azure/MS-AMP}{aka.ms/MS.AMP}.