Source-linked AI summary
Pretraining Large Language Models with NVFP4
NVIDIA, Felix Abecassis, Anjulie Agrusa, Dong Ahn, Jonah Alben, Stefania Alborghetti, Michael Andersch, Sivakumar Arayandi, Alexis Bjorlin, Aaron Blakeman, Evan Briones, Ian Buck, Bryan Catanzaro, Muya Chang, Jinhang Choi, Mike Chrzanowski, Eric Chung, Victor Cui, Steve Dai, Bita Darvish Rouhani, Carlo del Mundo, Deena Donia, Burc Eryilmaz, Henry Estela, Abhinav Goel, Oleg Goncharov, Yugi Guvvala, Robert Hesse, Russell Hewett, Herbert Hum, Ujval Kapasi, Brucek Khailany, Mikail Khona, Nick Knight, Alex Kondratenko, Ronny Krashinsky, Ben Lanir, Simon Layton, Michael Lightstone, Daniel Lo, Paulius Micikevicius, Asit Mishra, Tim Moon, Deepak Narayanan, Chao Ni, Abhijit Paithankar, Satish Pasumarthi, Ankit Patel, Mostofa Patwary, Ashwin Poojary, Gargi Prasad, Sweta Priyadarshi, Yigong Qin, Xiaowei Ren, Oleg Rybakov, Charbel Sakr, Sanjeev Satheesh, Stas Sergienko, Pasha Shamis, Kirthi Shankar, Nishant Sharma, Mohammad Shoeybi, Michael Siu, Misha Smelyanskiy, Darko Stosic, Dusan Stosic, Bor-Yiing Su, Frank Sun, Nima Tajbakhsh, Shelby Thomas, Przemek Tredak, Evgeny Tsykunov, Gandhi Vaithilingam, Aditya Vavre, Rangharajan Venkatesan, Roger Waleffe, Qiyu Wan, Hexin Wang, Mengdi Wang, Lizzie Wei, Hao Wu, Evan Wu, Keith Wyss, Ning Xu, Jinze Xue, Charlene Yang, Yujia Zhai, Ruoxi Zhang, Jingyang Zhu, Zhongbo Zhu
TL;DR
Large-scale FP4 pretraining could reduce the computational and memory demands of LLM training, but stability and convergence remain challenging. The paper develops a methodology around NVFP4 and validates it on a 12B model trained for 10T tokens, closely matching FP8 loss and downstream accuracy. The results establish public evidence of sustained 4-bit pretraining at multi-trillion-token scale, while leaving broader layer coverage and larger-scale validation for future work.
Problem
FP4 training promises lower computational and memory costs, but stable convergence remains challenging for large models over long token horizons.
Method
The paper combines NVFP4 with selective high precision, Random Hadamard transforms, 2D weight scaling, and stochastic gradient rounding.
Results
A 12B model trained on 10T tokens closely tracks FP8 loss and downstream accuracy, establishing sustained 4-bit pretraining at multi-trillion-token scale.
Takeaways & Limitations
NVFP4 can support stable and accurate large-scale pretraining when paired with the proposed methodology.
Takeaways & Limitations
The study leaves quantization of all linear layers, reduced high-precision use, and evaluation on larger models and other architectures for future work.
Abstract
from arXiv · showhide
Large Language Models (LLMs) today are powerful problem solvers across many domains, and they continue to get stronger as they scale in model size, training set size, and training set quality, as shown by extensive research and experimentation across the industry. Training a frontier model today requires on the order of tens to hundreds of yottaflops, which is a massive investment of time, compute, and energy. Improving pretraining efficiency is therefore essential to enable the next generation of even more capable LLMs. While 8-bit floating point (FP8) training is now widely adopted, transitioning to even narrower precision, such as 4-bit floating point (FP4), could unlock additional improvements in computational speed and resource utilization. However, quantization at this level poses challenges to training stability, convergence, and implementation, notably for large-scale models trained on long token horizons. In this study, we introduce a novel approach for stable and accurate training of large language models (LLMs) using the NVFP4 format. Our method integrates Random Hadamard transforms (RHT) to bound block-level outliers, employs a two-dimensional quantization scheme for consistent representations across both the forward and backward passes, utilizes stochastic rounding for unbiased gradient estimation, and incorporates selective high-precision layers. We validate our approach by training a 12-billion-parameter model on 10 trillion tokens -- the longest publicly documented training run in 4-bit precision to date. Our results show that the model trained with our NVFP4-based pretraining technique achieves training loss and downstream task accuracies comparable to an FP8 baseline. These findings highlight that NVFP4, when combined with our training approach, represents a major step forward in narrow-precision LLM training algorithms.
NVIDIA
Transformer Engine supports NVFP4 training.
- Transformer Engine provides support for NVFP4 training.
1. Introduction
The report motivates FP4 training as a way to reduce the computational and memory demands of large-scale LLM pretraining, while addressing its stability challenges with a targeted methodology. It validates the approach on a 12B-parameter model trained for 10T tokens against FP8.
- FP4 can provide a two- to three-fold arithmetic-performance boost and halve memory usage versus FP8.
- NVFP4 uses smaller micro-blocks and fractional-precision scale factors to capture local dynamic range more accurately than MXFP4.
- The methodology combines high-precision sensitive layers, 2D weight scaling, Random Hadamard transforms, and stochastic gradient rounding.
- Each methodology component is especially important for 4-bit training in large models and over long token horizons.
- A 12-billion-parameter LLM trained on 10 trillion tokens closely matches FP8 loss curves and downstream-task accuracies.
2. NVFP4 Format
NVFP4 improves on microscaling formats by using finer blocks, more precise block scales, and a tensor-level FP32 scale. These choices preserve dynamic-range information while retaining FP4 storage for most values.
- MX formats use block-wise shared scale factors to balance the limited range and precision of narrow floating-point formats.
- Quantization maps each block’s absolute maximum to the FP4 maximum, then decodes values using the reciprocal scale.
- NVFP4 reduces block size from 32 to 16 elements and stores block scales in E4M3 rather than UE8M0.
- At least 6.25% of each NVFP4 block is encoded at near-FP8 precision, while the remaining values use FP4.
- A tensor-level FP32 scale adjusts the distribution so block scales remain representable in E4M3.
- NVFP4’s smaller blocks and more precise scaling improve outlier accuracy and reduce small values quantized to zero.
3. Training with NVFP4
NVFP4 pretraining remains stable across 10T tokens, with validation loss and downstream accuracy closely tracking FP8. A small amount of higher-precision training can further reduce the late-stage loss gap.
- A 12B hybrid Mamba-Transformer model was trained for 10T tokens with NVFP4 against an FP8 reference.
- NVFP4 validation loss stays within 1% relative error during stable training and rises slightly above 1.5% late in training.
- NVFP4 matches FP8 across downstream domains including reasoning, mathematics, coding, and commonsense tasks, with a slight coding shortfall.
- Changing from NVFP4 to BF16 or potentially MXFP8 during decay mitigates the late-stage loss gap.
- Stable long-horizon training preserves accuracy relative to higher-precision baselines and supports scalable 4-bit training.
4. Training Methodology
The methodology combines mixed precision, outlier management, consistent two-dimensional scaling, and stochastic rounding to stabilize 4-bit training over long horizons. Ablations indicate that removing any major component worsens convergence, while selected layers and tensors remain in higher precision.
- Method overview: The approach combines higher-precision sensitive layers, Random Hadamard transforms, 2D weight scaling, and stochastic rounding for effective 4-bit training.The recommended configuration also uses 1D scaling for activations and gradients, with round-to-nearest-even for weights and activations.
- Ablations: Removing any methodology component leads to worse convergence for the 12B model trained over 10T tokens.The ablations remove stochastic rounding, Random Hadamard transforms, two-dimensional scaling, or high-precision blocks one at a time.
- Mixed precision: Most linear-layer GEMMs use FP4 inputs and produce outputs in BF16 or FP32 across forward and backward computation.The three GEMMs are Fprop, Dgrad, and Wgrad; embeddings, normalization, nonlinearities, attention components, optimizer states, and key weights or gradients remain higher precision.
- Mixed precision: Sensitive final linear layers require more dynamic range and mantissa than FP4 provides, so a small fraction remains in BF16 or MXFP8.For the 12B model, the conservative configuration keeps the first two and final eight blocks in BF16, representing 16% of linear layers.
- Random Hadamard transforms: Random Hadamard transforms redistribute outliers, with empirical benefit observed for Wgrad inputs but no measurable benefit for Fprop and Dgrad.Randomization benefits larger models trained over longer token horizons, while small-scale training remains stable with the standard transform.
- Two-dimensional scaling: Two-dimensional scaling preserves consistent quantized representations across forward and backward passes, addressing a mismatch caused by transposed backward computation.Without this consistency, the same tensor can receive distinct quantized representations, effectively breaking the chain rule.
- Stochastic rounding: Stochastic rounding is essential for gradient convergence in the 12B model, whereas applying it to forward tensors increases quantization error.Other backward-pass tensors do not benefit from stochastic rounding, indicating that gradients are the primary source of rounding bias.
5. NVFP4 and MXFP4
The study compares NVFP4 and MXFP4 using an 8B hybrid Mamba-Transformer model trained for 1T tokens. NVFP4 converges to lower loss, while MXFP4 requires substantially more tokens to match it.
- Model and setup: The comparison uses an 8B hybrid Mamba-Transformer trained on 1T tokens with the same dataset as the 12B model.Training uses two phases of data blending between the first 60% and last 40% of training.
- Model and setup: Both FP4 settings retain the last eight blocks in BF16, comprising about 15% of the model.The reference model is pretrained in BF16, while FP4 pretraining uses the methodology described for MXFP4 and NVFP4.
- Results: NVFP4 pretraining converges to better loss than MXFP4, with relative error around 1.5% for NVFP4 versus 2.5% for MXFP4.Figure 6 compares training-loss difference and validation perplexity across token budgets.
- Results: MXFP4 matches NVFP4 loss only after 36% more tokens, using 1.36T instead of 1T tokens.The additional-token comparison indicates a considerable increase in training time for MXFP4.
6. Conclusions
The report finds that NVFP4 can support stable, accurate large-scale pretraining when paired with its targeted methodology. A 12B model trained on 10T tokens closely tracks FP8 in loss and downstream accuracy, while broader validation remains future work.
- Conclusions: NVFP4 pretraining is stable and accurate with 2D weight scaling, Random Hadamard transforms, stochastic rounding, and selective higher-precision layers.The methodology is designed to improve training stability and convergence.
- Conclusions: A 12B hybrid Mamba-Transformer trained on 10T tokens closely tracks the FP8 baseline in loss and downstream accuracy.The report describes this as the first public evidence of sustained 4-bit pretraining at multi-trillion-token scale.
- Conclusions: NVFP4 reaches comparable loss with fewer tokens than MXFP4, indicating efficiency gains without sacrificing accuracy.The comparison provides an initial view of memory and compute efficiency benefits and convergence trade-offs among FP4 formats.
- Future work: Future work includes quantizing all linear layers, reducing remaining high-precision layers, extending NVFP4 to attention and communication paths, and evaluating larger models and additional architectures.The report also identifies post-training scenarios and longer token horizons for further evaluation.
A.2. 8B hybrid Mamba-Transformer
The 8B hybrid Mamba-Transformer uses 52 blocks combining Self-Attention, FFN, and Mamba-2 components. Its NVFP4 configuration leaves the final eight blocks in BF16 while applying the described quantization procedure elsewhere.
- Model architecture: The 8B model has 52 blocks: 4 Self-Attention, 24 FFNs, and 24 Mamba-2 blocks.
- Model architecture: Its hidden dimension is 4096, FFN hidden dimension is 21504, and Grouped-Query Attention uses 32 query heads and 4 key-value heads.
- Model architecture: Mamba-2 blocks use 8 groups, state dimension 128, head dimension 64, expansion factor 2, and convolution window size 4.
- Hyperparameters: The model is trained on 1 trillion tokens with sequence length 8192, batch size 768, and a learning rate decaying from 8.0 · 10^-4 to 8.0 · 10^-6 over the last 15% of training.
- Precisions: All linear layers use NVFP4 except those in the last eight FFN or Mamba-2 blocks, which remain in BF16.
- Quantization: NVFP4 conversion uses global FP32 tensor-level scaling followed by local block-level scaling, with decode scales stored for the matrix multiplication.
B.4. Remarks on MXFP4 and NVFP4 scale factor
NVFP4 improves on MXFP4 scaling by using more precise block scales and Random Hadamard transforms. These choices better use FP4’s representable range while dispersing outliers before quantization.
- MXFP4 limitation: MXFP4 power-of-two scale factors can waste FP4 samples and reduce the utilized dynamic range.
- NVFP4 scaling: NVFP4 uses a precise E4M3 block scale that maps each block’s amax closer to the FP4 maximum representable value.
- NVFP4 scaling: This scaling maximizes FP4 sample utilization and preserves more of FP4’s dynamic range.
- Random Hadamard transforms: Random Hadamard transforms rotate tensors before quantization, computing the scale factor in the rotated space.
- Random Hadamard transforms: The transforms can be applied to both matrix-multiplication operands because their orthogonal factors cancel through H H^T = I.
- Implementation: Tiled Hadamard transforms have cost m k d multiply-adds and can be fused with other layers to reduce device-memory round-trips.
D. Switching to Higher Precision
Switching selected tensors from NVFP4 to higher precision near the end of training can recover much of the loss gap. Forward-pass switching is especially effective while limiting the computation affected.
- Loss recovery: 8.2T tokens is the reported switch point where validation loss matches the FP8 baseline, while switching at 10T tokens is slightly worse.
- Loss recovery: Switching to higher precision shortly before learning-rate decay is recommended for full loss recovery, while switching at the end provides notable improvement with minimal runtime effect.
- Forward versus backward pass: Most of the FP4 loss gap arises from quantizing forward-pass tensors.
- Forward versus backward pass: Forward-pass switching from 8.2T tokens reduces the 12B model’s relative error from 1.5% to 0.5%, whereas backward-pass switching shows no improvement in these models.
- Forward versus backward pass: Focusing on the forward pass limits overhead because only about 6% of total computations are switched.
E. Ablation of Training Methodology
The ablation study examines how combining NVFP4 techniques affects validation loss in a 1.2B model. Its base method quantizes all layers and tensors with standard NVFP4 scaling and round-to-nearest-even.
- Ablation setup: Figure 8 combines BF16 linear layers in the last four blocks, 2D weight scaling, Random Hadamard transforms on Wgrad, and stochastic rounding on gradients.
- Evaluation: The combined-technique ablation is evaluated using relative validation-loss differences for a 1.2B model trained on 1T tokens.
- Ablation setup: The base method quantizes all layers to NVFP4, applies standard NVFP4 scaling to all tensors, and uses round-to-nearest-even throughout.
- Base method: The base method causes models to diverge early in training.
E.2. Layer sensitivity
FP4 sensitivity is concentrated in specific linear layers and training operations. Stable NVFP4 training depends on preserving sensitive final layers, applying stochastic rounding selectively, and choosing RHT and scaling schemes that maintain effective representations.
- Layer sensitivity: Training diverges when every linear layer uses FP4, but retaining the last four blocks in BF16 restores loss convergence.Keeping only early blocks in higher precision does not improve stability unless combined with higher-precision final blocks.
- Layer sensitivity: The final layers have larger Wgrad quantization errors, making quantization-error metrics a possible guide for selective higher-precision retention.
- Stochastic rounding: Stochastic rounding stabilizes convergence when applied to gradients, whereas applying it to activations or weights causes divergence.For the 12B model, gradients entering both Dgrad and Wgrad require stochastic rounding for proper convergence.
- Random Hadamard transforms: Applying RHT to Wgrad inputs improves validation loss, while applying it to Fprop or Dgrad inputs degrades model quality.The authors hypothesize that RHT-induced quantization error can offset its outlier-removal benefit for some GEMMs.
- Random Hadamard transforms: 16 × 16 Hadamard transforms provide a cost-effective choice, while 4 × 4 increases loss and 128 × 128 yields a minor quality benefit.The results support 16 × 16 transforms without compromising model accuracy.
- Tensor consistency: Different quantized weight representations across forward and backward passes worsen loss, while 2D block scaling improves it despite larger block granularity.Activations are less sensitive to representation consistency and are affected mainly during later learning-rate decay.