面向 Vision Transformer 的振荡抑制 MXFP4 训练
机器学习
2025-07-10 v2 人工智能
计算机视觉与模式识别
摘要
以 FP4 精度预训练 Transformer 正成为一种有望获得显著加速的方法，但其伴随着相当大的精度损失。微缩放（MX）数据格式提供了一种细粒度的逐组量化方法，以提高 FP4 格式的表示能力，并受到下一代 Blackwell GPU 架构的支持。然而，使用 MXFP4 数据格式进行训练仍会导致显著的性能退化，且对其原因缺乏系统研究。在这项工作中，我们提出了一种用于更精确 FP4 训练的新型训练方法 TetraJet。我们全面评估了训练中涉及的所有量化器，并确定前向传播中的权重振荡问题是 MXFP4 训练退化的主要来源。因此，我们引入了两种新方法：EMA 量化器和自适应爬坡优化器，以解决振荡问题。在 Vision Transformer 上的大量实验表明，TetraJet 始终优于现有的 4 位训练方法，且 Q-EMA 和 Q-Ramping 能够通过有效减少振荡提供额外的性能提升。与基线相比，我们将精度退化降低了 50% 以上，甚至能够获得与全精度训练相当的性能。代码可在 https://github.com/thu-ml/TetraJet-MXFP4Training 获取
引用
@article{arxiv.2502.20853,
title = {Oscillation-Reduced MXFP4 Training for Vision Transformers},
author = {Yuxiang Chen and Haocheng Xi and Jun Zhu and Jianfei Chen},
journal= {arXiv preprint arXiv:2502.20853},
year = {2025}
}