E5M3 块缩放实现稳定的 FP4 预训练
来自 arXiv · 期号 2026-09-03
摘要
4 位浮点（FP4）预训练很难稳定，因为 E2M1 尾数能表达的量级范围太窄。 NVIDIA 的 Transformer Engine NV4 方案用逐张量缩放、随机化 Hadamard 变换（RHT）和 BF16 最后几层来弥补，但这些操作都跑在 FP4 矩阵乘法之外。 这篇论文把 E2M1 尾数换成无符号 E5M3 块缩放。E5M3 的量程更宽，允许周期性张量缩放；方案对反向梯度做选择性随机舍入，去掉 RHT，所有可用的内部线性层都用 FP4。 用 Nemotron-H 8B 模型预训练了近 1900 亿 token。与 Transformer Engine NV4 相比，块-16 方案的最终窗口训练损失更低，验证损失（保留集负对数似然）也更低。三个聚合指标上的量化推理下游分数都更高。 去掉 RHT 和 BF16 最后块豁免的消融实验，把模型体 token 吞吐量提升了 21.2%。 这些结果表明，用更简单的方案就能实现端到端软件模拟的 FP4 预训练，也为原生支持 E5M3 块缩放提供了动机。