🤖 AI Summary
This work addresses the limitations of existing 4-bit pretraining methods, which quantize only the linear layers of Transformers while neglecting optimizer states, computational stability, and attention mechanisms. The paper proposes NVFP4, the first end-to-end full-stack 4-bit pretraining framework for large language models, featuring a modular precision strategy that jointly optimizes three core components: LoRA-SVD denoising to stabilize weight updates, robust second-moment storage in AdamW combined with Newton–Schulz iterations to preserve optimizer numerical accuracy, and mixed-precision quantization of Q/K/V activations and gradients with unified tensor reuse to enhance attention stability. Evaluated on a 3B-parameter model trained on 64B tokens, NVFP4 incurs only a 1.47% loss gap relative to BF16 training, demonstrating—for the first time—the feasibility and efficiency of stable, full-stack 4-bit LLM pretraining.
📝 Abstract
Recent NVFP4 pretraining methods mainly target transformer linear layers, leaving optimizer states, optimizer arithmetic and attention underexplored in 4-bit pipelines. This critical gap blocks stable full-stack 4-bit pretraining, as the three core modules exhibit unique numerical failure patterns: linear layers hit hard quantization noise limits with dimension-propagated error amplification; AdamW second moments are heavy-tailed non-negative values fragile to low-precision denominators; attention carries error-prone computation paths demanding strict forward-backward quantization consistency. We propose Full-Stack FP4, the first complete NVFP4 pretraining framework resolving all three stability bottlenecks via module-wise precision strategies. For linear projections, LoRA-SVD lightweight decomposition suppresses quantization noise and breaks the direct-quantization error ceiling, shrinking the linear-only loss gap from 1.40% to 0.61%. For optimizers, we design AdamW second-moment transformation for robust NVFP4 storage and fully support native NVFP4 Newton-Schulz iterations for the Root (Muon) optimizer. For attention, a mixed-precision scheme quantizes Q/K/V and backward dS while guarding vulnerable paths in BF16, paired with unified tensor reuse to sustain forward-backward alignment. We further analyze fast error accumulation in naive low-bit matrix multiplication and the extreme sensitivity of PV / dOV^T attention branches. All modules are plug-and-play with cumulative stability and efficiency improvements. Our 3B/64B-token pretraining validates near-BF16 performance with merely 1.47% loss gap, verifying feasible stable end-to-end NVFP4 LLM pretraining.