Scaling FP8 training to trillion-token LLMs

📄 arXiv: 2409.12517v2 📥 PDF

作者: Maxim Fishman, Brian Chmiel, Ron Banner, Daniel Soudry

分类: cs.LG, cs.AI

发布日期: 2024-09-19 (更新: 2025-02-10)

🔗 代码/项目: GITHUB

💡 一句话要点

提出Smooth-SwiGLU，解决FP8训练LLM时因SwiGLU激活函数引起的长期不稳定性问题。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: FP8训练 大型语言模型 SwiGLU激活函数 Smooth-SwiGLU Adam优化器 量化训练 Intel Gaudi2 训练稳定性

📋 核心要点

  1. 现有FP8训练大型语言模型的方法在长时间训练中存在不稳定性，限制了模型规模和训练数据量的扩展。
  2. 提出Smooth-SwiGLU激活函数，并对Adam优化器矩进行FP8量化，以解决FP8训练中的长期不稳定性问题。
  3. 在256个Intel Gaudi2加速器上，使用FP8精度训练了一个7B参数模型，吞吐量相比BF16基线提升约34%，性能相当。

📝 摘要（中文）

本文首次使用FP8精度在高达2万亿token的数据集上训练大型语言模型，相比之前的限制提高了20倍。通过这些扩展的训练，我们发现了FP8训练中存在的关键不稳定性，这些不稳定性在早期较短的训练中无法观察到。我们发现这些不稳定性是由SwiGLU激活函数对异常值的放大引起的。有趣的是，我们通过分析和实验表明，这种放大仅发生在长时间的训练中，并将其与SwiGLU权重对齐过程联系起来。为了解决这个新发现的问题，我们引入了Smooth-SwiGLU，这是一种新颖的修改，可确保稳定的FP8训练，而不会改变函数行为。我们还首次展示了Adam优化器矩的FP8量化。结合这些创新，我们成功地在256个Intel Gaudi2加速器上使用FP8精度训练了一个7B参数模型，实现了与BF16基线相当的结果，同时提供了高达约34％的吞吐量提升。参考实现可在https://github.com/Anonymous1252022/Megatron-DeepSpeed获得。

🔬 方法详解

问题定义：论文旨在解决使用FP8精度训练大规模语言模型时遇到的长期不稳定性问题。现有方法在训练数据量和模型规模扩展时，由于FP8的精度限制，容易出现训练崩溃，尤其是在长时间训练的情况下。之前的研究未能充分揭示这些长期不稳定性，限制了FP8在LLM训练中的应用。

核心思路：论文的核心思路是识别并解决导致FP8训练长期不稳定性的根本原因。通过深入分析，作者发现SwiGLU激活函数在长时间训练中会放大异常值，导致梯度爆炸和训练崩溃。因此，论文提出了一种改进的激活函数Smooth-SwiGLU，以抑制异常值放大，从而稳定FP8训练。

技术框架：论文的技术框架主要包括以下几个部分：1) 使用FP8精度训练LLM；2) 识别SwiGLU激活函数引起的长期不稳定性；3) 提出Smooth-SwiGLU激活函数；4) 对Adam优化器矩进行FP8量化；5) 在硬件加速器上进行实验验证。整体流程是从问题发现到解决方案提出，再到实验验证的闭环。

关键创新：论文最重要的技术创新点在于Smooth-SwiGLU激活函数的提出。该函数通过修改SwiGLU的计算方式，有效地抑制了异常值的放大，从而稳定了FP8训练。与原始SwiGLU相比，Smooth-SwiGLU在保持函数行为的同时，避免了长时间训练中的梯度爆炸问题。此外，对Adam优化器矩进行FP8量化也是一个重要的创新，进一步提升了FP8训练的效率。

关键设计：Smooth-SwiGLU的具体实现细节未知，但其核心思想是在SwiGLU的计算过程中引入某种平滑机制，以限制输出值的范围，防止异常值被过度放大。论文还对Adam优化器矩进行了FP8量化，具体量化方法未知，但其目的是减少内存占用和计算量，从而提升训练效率。损失函数和网络结构与标准的LLM训练流程保持一致，没有进行显著修改。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，使用Smooth-SwiGLU和FP8量化的Adam优化器，可以在256个Intel Gaudi2加速器上成功训练一个7B参数的语言模型，达到与BF16基线相当的性能，同时吞吐量提升约34%。这证明了Smooth-SwiGLU的有效性和FP8训练在LLM中的潜力。

🎯 应用场景

该研究成果可应用于大规模语言模型的训练，尤其是在资源受限的硬件平台上。通过使用FP8精度，可以在保证模型性能的同时，显著降低内存占用和计算成本，从而加速模型训练和部署。这对于推动人工智能在各个领域的应用具有重要意义，例如自然语言处理、机器翻译、文本生成等。

📄 摘要（原文）

We train, for the first time, large language models using FP8 precision on datasets up to 2 trillion tokens -- a 20-fold increase over previous limits. Through these extended training runs, we uncover critical instabilities in FP8 training that were not observable in earlier works with shorter durations. We trace these instabilities to outlier amplification by the SwiGLU activation function. Interestingly, we show, both analytically and empirically, that this amplification happens only over prolonged training periods, and link it to a SwiGLU weight alignment process. To address this newly identified issue, we introduce Smooth-SwiGLU, a novel modification that ensures stable FP8 training without altering function behavior. We also demonstrate, for the first time, FP8 quantization of both Adam optimizer moments. Combining these innovations, we successfully train a 7B parameter model using FP8 precision on 256 Intel Gaudi2 accelerators, achieving on-par results with the BF16 baseline while delivering up to a $\sim 34 \%$ throughput improvement. A reference implementation is supplied in https://github.com/Anonymous1252022/Megatron-DeepSpeed.