蚂蚁:UFP4修正FP4训练收缩偏差

📖标题:Rethinking Shrinkage Bias in LLM FP4 Pretraining: Geometric Origin, Systemic Impact, and UFP4 Recipe
🌐来源:arXiv, 2606.20381v1
🛎️文章简介
🔸研究问题:如何解决主流E2M1格式在FP4预训练中因几何不对称导致的系统性收缩偏差及训练不稳定问题?
🔸主要贡献:论文揭示了非均匀格式的收缩偏差根源,提出基于均匀网格的UFP4配方,实现了全路径RHT并显著提升训练精度。
📝重点思路
🔸识别收缩偏差几何根源:指出E2M1等非均匀格式因RTNE量化桶的几何不对称性,产生系统性负向舍入误差,而E1M2等均匀网格可从根本上规避此误差。
🔸揭示偏差累积与RHT恶化机制:证明收缩偏差在网络层间呈乘法累积导致信号衰减,且随机哈达玛变换(RHT)会将张量质量转移至E2M1最不对称的区间,反而加剧训练不稳定性。
🔸设计UFP4均匀训练配方:采用E1M2/INT4风格均匀网格,将RHT安全扩展至前向、数据梯度及权重梯度全部三个GEMM路径,仅对上游梯度dY使用随机舍入。
🔸验证硬件融合效率:实现RHT与FP4量化的算子融合,在SM90/SM100架构上仅增加约6%-7%延迟,证明了全路径RHT在工业级训练中的工程可行性。
🔎分析总结
🔸RHT改变最优格式选择:实验证实RHT将张量从动态范围受限转为局部分辨率受限,使均匀网格E1M2在旋转后的信噪比和桶利用率上全面超越E2M1。
🔸长程训练损失显著降低:在Dense 1.5B、MoE 7.9B及MoE 124B模型上,UFP4相比E2M1基线持续保持更低的BF16相对损失退化,验证了均匀网格的有效性。
🔸缩放定律优势持久:通过10M至324M参数的缩放实验,证实E1M2的损失曲线始终低于E2M1,且FP4与BF16的差距随计算量增加未出现恶化趋势。
🔸全路径RHT不可或缺:消融实验表明,仅在均匀网格下启用覆盖全部三个GEMM的全路径RHT才能获得最佳损失,限制了RHT范围会导致性能明显下降。
🔸E2M1无法模拟均匀网格:尝试通过限制E2M1数值范围来近似均匀网格的方案均导致性能劣化,说明原生均匀格式支持是必要的。
💡个人观点
论文从量化网格几何本质出发,发现RHT与E2M1存在根本性冲突,出以均匀网格替代非均匀网格作为FP4训练原语。
更多推荐



所有评论(0)