UE5M3 FP4 Block Scaling for Stable Language Model Pretraining
Introduces UE5M3 block scaling with periodic tensor scaling, enabling stable FP4 pretraining of a 8B model with 190B tokens, outperforming NVFP4 in loss metrics.
Key Findings
Methodology
This paper proposes a UE5M3 unsigned block scaling scheme combined with periodic tensor maximum refresh and selective stochastic rounding, omitting RHT transforms and BF16 final layers. The approach employs a simplified recipe for end-to-end FP4 pretraining, training an 8B parameter Nemotron-H model over nearly 190 billion tokens. The method involves controlling scale saturation through a wider block scale range (T from 448 to 2048), using periodic sampling of tensor maxima every 50 steps, and applying stochastic rounding only to upstream gradients. All internal linear layers utilize FP4, reducing complexity. Comparative experiments demonstrate lower training loss and validation negative log-likelihood than NVIDIA’s NVFP4, with increased throughput after ablation of RHT and BF16 layers.
Key Results
- The proposed scheme achieves lower final-window training loss and validation negative log-likelihood compared to NVFP4, with the 8B model trained on 190B tokens showing superior performance. Expanding block scale range (T from 448 to 2048) improves small gradient representation, reducing underflow and saturation events. Ablation removing RHT and BF16 final layers increases token throughput by 21.2%. The model’s downstream point estimates are higher across three reported metrics, indicating better quantized inference accuracy.
- Experiments confirm that wider block scales (T=2048) allow small gradients to be represented more accurately, leading to more stable training. The method’s simplicity, with fewer auxiliary operations, reduces implementation complexity while maintaining or improving model quality. The ablation studies validate the importance of scale range extension and selective stochastic rounding.
- The approach demonstrates that a wider, unsigned block-scale format paired with periodic tensor maximum updates can stabilize FP4 training without RHT or BF16 layers, paving the way for native hardware support and more efficient low-bit quantized training pipelines.
Significance
This work addresses the longstanding challenge of stable FP4 pretraining by expanding the dynamic range through UE5M3 block scales and periodic tensor scaling. It simplifies the training pipeline, reducing computational overhead and implementation complexity, while achieving better model performance. The findings have significant implications for deploying large language models efficiently, especially in hardware-constrained environments. By demonstrating a software-emulated end-to-end training process that surpasses previous methods, the research provides a strong foundation for future hardware support of UE5M3 block scaling, potentially enabling faster, more energy-efficient training of billion-parameter models.
Technical Contribution
The core technical innovation lies in reusing the sign bit for an additional exponent bit in UE5M3, vastly increasing the scale range (up to 61,440) compared to traditional E4M3 (448). This wider range allows delayed tensor scaling without RHT, simplifying the training recipe. The method employs periodic maximum value sampling every 50 steps, with a target scale T adjustable between 448 and 2048, to prevent saturation. Selective stochastic rounding is applied only to upstream gradients, reducing bias. All internal linear layers operate in FP4, with the entire process simulated via software. This approach outperforms NVFP4 in both training loss and inference accuracy, demonstrating a new paradigm for low-bit quantized training.
Novelty
This study is the first to systematically incorporate UE5M3 unsigned block scales with periodic tensor maximum refresh in end-to-end FP4 pretraining. It eliminates the need for RHT and BF16 final layers, simplifying the pipeline while maintaining or improving model performance. The wider scale range (up to 61,440) significantly enhances the stability of FP4 training, especially for small gradients, representing a major step forward in low-bit quantization research. The combination of scale range extension, periodic refresh, and selective stochastic rounding forms a novel, effective recipe that surpasses existing NVFP4 methods.
Limitations
- The current implementation relies on software emulation, lacking native hardware support, which may limit real-world performance gains. Hardware integration remains future work.
- Scale target T requires manual tuning; improper settings can lead to saturation or underflow, affecting training stability.
- Extreme scale variations or highly sensitive tasks might still encounter saturation or loss of small gradient information, necessitating further adaptive scale mechanisms.
Future Work
Future directions include integrating native hardware support for UE5M3 block scales, automating scale target tuning, and exploring adaptive scale adjustment strategies. Extending the approach to larger models and diverse tasks will test its robustness and generality. Additionally, combining UE5M3 with other quantization schemes and hardware-aware optimizations could further improve training efficiency and model accuracy, accelerating the deployment of low-bit quantized large language models.
AI Executive Summary
The rapid growth of large language models has driven the need for more efficient training methods, especially in low-precision formats like FP4. While FP8 and INT4 have seen hardware support, FP4 remains challenging due to its narrow dynamic range, which hampers training stability. Existing solutions such as NVIDIA’s Transformer Engine incorporate complex mechanisms like randomized Hadamard transforms (RHT), two-dimensional weight scaling, and BF16 final layers to stabilize training, but these add significant complexity.
In response, this study introduces a simplified yet effective approach based on UE5M3 unsigned block scaling combined with periodic tensor maximum refresh. By extending the block scale range (T from 448 to 2048), the method allows delayed tensor scaling without RHT, reducing computational overhead. Selective stochastic rounding is applied only to upstream gradients, further simplifying the pipeline. The entire training process is software-emulated, demonstrating that a 8B parameter Nemotron-H model can be trained on nearly 190 billion tokens with lower training loss and better validation metrics than NVFP4.
Experimental results show that wider block scales improve the representation of small gradients, reducing saturation and underflow issues. The ablation studies reveal a 21.2% increase in token throughput when removing RHT and BF16 layers, confirming the efficiency gains. The approach’s simplicity, combined with its performance, suggests promising avenues for native hardware support, potentially enabling faster, more energy-efficient training of billion-parameter models.
Overall, this work advances low-bit quantization techniques by expanding dynamic range through innovative block scaling, paving the way for practical large-scale language model training with reduced computational costs and complexity. It sets a new benchmark for FP4 pretraining, with broad implications for both academia and industry in deploying large models efficiently.
Deep Dive
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.