TetraJet-v2: Accurate NVFP4 Training for Large Language Models with Oscillation Suppression and O...
文章总结与翻译一、主要内容本文针对大型语言模型(LLMs)低精度全量化训练(FQT)中4位格式(如NVFP4)存在的性能损失问题,提出了TetraJet-v2——一种端到端的4位FQT方法。该方法在所有线性层的激活、权重和梯度中均采用NVFP4格式,核心解决了低精度训练中的两个关键问题:权重振荡和异常值影响。通过三大核心技术实现高效训练:1)无偏双块量化方法,适配NVFP4的数值需求并保证梯度估计无偏;2)OsciReset算法,通过识别振荡权重并将其重置到量化区间中心,抑制权重振荡;3)OutControl算法,通过静态选择异常值通道并保留高精度,同时在反向传播中结合随机哈达玛变换(RHT),兼顾前向和反向的异常值精度。实验在OLMo-2模型(70M、150M、370M参数)和OLMo-2-Mix-1124数据集(50B-200B tokens)上验证,TetraJet-v2-full版本相比现有FP4训练方法,将与全精度训练的性能差距平均缩小51.3%,在下游任务(如ARC、MMLU等)中也取得最优平均性能。此外,文章还通过消融实验揭示了前向传播和MLP.ffn2层对量化最敏感,且FP6×FP4格式有望进一步提升性能。二、创新点无偏双块量化设计:针对NVFP4的E4M3缩放因子数值范围限制,引入1×128外块+1×16内块的双块量化,既满足数值需求又提升硬件友好性,同时通过随机舍入保证梯度估计无偏。首个LLM低精度训练权重振荡解决方案:提出OsciRese