将 FP8 训练扩展至万亿 token 的大语言模型
机器学习
2025-02-11 v2 人工智能
摘要
我们首次在数据量达 2 万亿 token 的大规模语言模型上使用 FP8 精度进行训练 —— 相较于前期限制提升了 20 倍。通过这些长时间的训练运行，我们发现了在以往时长较短的工作中未观察到的 FP8 训练中的关键不稳定性。我们追溯这些不稳定性到 SwiGLU 激活函数对异常值的放大。有趣的是，我们既通过分析又通过实证研究表明，这种放大仅在长期训练期间才会发生，并与 SwiGLU 权重对齐过程相关。为解决此新发现的问题，我们引入 Smooth-SwiGLU，这是一种 novel modification，确保在不改变函数行为的前提下实现稳定的 FP8 训练。我们还首次对 Adam 优化器的两个时序进行 FP8 量化。结合这些创新措施，我们成功在 256 个 Intel Gaudi2 加速器上使用 FP8 精度训练一个 7B 参数模型，实现与 BF16 基线相当的效果，同时提升约 34% 的吞吐量。参考实现已发布于 https://github.com/Anonymous1252022/Megatron-DeepSpeed。
引用
@article{arxiv.2409.12517,
title = {Scaling FP8 training to trillion-token LLMs},
author = {Maxim Fishman and Brian Chmiel and Ron Banner and Daniel Soudry},
journal= {arXiv preprint arXiv:2409.12517},
year = {2025}
}