UE5M3 FP4 Block Scaling for Stable Language Model Pretraining

2026-09-02

AI Summary

本文针对4位浮点数预训练不稳定的问题，提出了一种结合E2M1载荷与无符号E5M3块尺度的方法，并通过选择性随机舍入和全FP4内部线性运算来实现更稳定的语言模型预训练。

Abstract

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.

Problem

• FP4
• pretraining
• stability
• E2M1 payload
• magnitude range

Innovation

• UE5M3 FP4
• block scaling
• selective stochastic rounding
• E2M1 payload
• E5M3