用更宽范围的E5M3块缩放,实现高效4比特语言模型预训练。
UE5M3 FP4 Block Scaling for Stable Language Model Pretraining

- 采用E5M3块缩放替代传统方法,支持周期性张量缩放。
- 在1900亿词上预训练80亿参数模型,训练与验证损失更低。
- 无需随机哈达玛变换,全链路使用FP4,吞吐提升21.2%。
4比特浮点(FP4)预训练因E2M1表示范围过窄而困难。NVIDIA Transformer Engine方案通过当前张量缩放、随机哈达玛变换(RHT)和最终层使用bfloat16(BF16)解决,但引入额外计算开销。本文改用无符号E5M3(p{})块缩放,其更广范围支持周期性张量缩放;提出的新方案采用选择性随机舍入处理反向梯度,省略RHT,并在所有可适配的内部线性层中使用FP4。在近1900亿词上预训练了Nemotron-H 8B模型。相比Transformer Engine v{},本文方法在各自量化推理策略下,最终窗口训练损失更低,验证损失(持有样本负对数似然)更低,下游任务点估计在三个指标上均更高。移除RHT与最终块不量化双重限制的原生执行消融实验,使模型主体词元吞吐率提升21.2%。结果表明,该方案实现了端到端软件模拟的p{}预训练,流程更简洁,推动对p{}块缩放的原生支持。
原文摘要 · Abstract (English)
Stable 4-bit floating-point (FP4) pretraining is difficult because the E2M1 payload represents only a narrow range of magnitudes. NVIDIA's Transformer Engine \nv{} recipe addresses this with current-tensor scaling, a randomized Hadamard transform (RHT), and bfloat16 (BF16) final layers, adding work outside the FP4 matrix multiplications. We instead pair E2M1 payloads with unsigned E5M3 (\ue{}) block scales. Their wider range permits periodic tensor scaling, while our recipe applies selective stochastic rounding to backward gradients, omits RHT, and uses FP4 in all eligible internal linears. We pretrain a Nemotron-H 8B model for nearly 190 billion tokens. Compared with Transformer Engine \nv{}, the proposed block-16 recipe finishes with lower final-window training loss and, under their respective quantized-inference policies, lower validation loss measured as held-out negative log-likelihood. Its quantized-inference downstream point estimates are also higher on all three reported aggregates. A native \nv{} execution ablation that jointly removes RHT and the BF16 final-block exemption increases measured model-body token throughput by 21.2\%. These results demonstrate end-to-end software-emulated \uefp{} pretraining with a simpler recipe and motivate native support for \ue{} block scaling.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。