arXiv:2409.12517cs.LGcs.AI2024-09ICLR被引 70

首次用FP8训练万亿级数据大模型,解决精度下降问题。

Scaling FP8 training to trillion-token LLMs

  • 引入Smooth-SwiGLU缓解长时训练中梯度异常放大。
  • 首次实现Adam优化器动量的FP8量化,支持2万亿token训练。
  • 70亿参数模型在34%更高吞吐下媲美16位基线性能。

我们首次在高达2万亿令牌的数据集上使用FP8精度训练大语言模型——较以往上限提升20倍。通过延长训练周期,发现此前未察觉的FP8训练不稳定性,其根源在于SwiGLU激活函数对异常值的放大效应。我们从理论和实证层面证明,该现象仅在长期训练中显现,并与SwiGLU权重对齐过程相关。为此提出Smooth-SwiGLU,可在不改变函数行为的前提下实现稳定训练。同时首次实现Adam优化器动量的FP8量化。结合这些创新,我们在256块Intel Gaudi2加速器上成功训练70亿参数模型,性能与BF16基线相当,吞吐最高提升约34%。参考实现已公开于https://github.com/Anonymous1252022/Megatron-DeepSpeed。

原文摘要 · Abstract (English)

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.

FP8大模型训练量化优化器

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。