论文精选73°

UE5M3 FP4块缩放实现稳定语言模型预训练

UE5M3 FP4 Block Scaling for Stable Language Model Pretraining

精选理由

NVIDIA提出新方法让FP4预训练更稳定,训练速度提升21.2%,模型性能全面超越现有方案。

AI 摘要

研究人员提出了一种新的UE5M3 FP4块缩放方法,用于稳定4位浮点预训练。该方法通过将E2M1负载与无符号E5M3块缩放配对,在Nemotron-H 8B模型上预训练了近1900亿个token。与Transformer Engine v方法相比,新方法实现了更低的最终窗口训练损失和验证损失,量化推理下游点估计在所有三个报告聚合上表现更好。

原文 · arXiv cs.LG

UE5M3 FP4 Block Scaling for Stable Language Model Pretraining

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.