Full-Stack FP4 Narrows Pretraining Loss Gap

Four-bit LLM training is no longer just a question of whether linear layers can tolerate quantization. A pretraining stack also has optimizer state, optimizer arithmetic, and attention paths whose numerical failures do not look the same as projection-layer noise. This paper argues that stable NVFP4 pretraining requires different precision choices for each module, rather than one blanket rule across the transformer.

Core Contribution

The authors introduce Full-Stack FP4, a complete NVFP4 pretraining recipe for transformer language models. The contribution is not a new model architecture; it is a systems-and-numerics design that pushes 4-bit arithmetic beyond the linear projections usually targeted by prior NVFP4 work. The paper separates the problem into three bottlenecks: dimension-amplified quantization error in linear projections, fragile low-precision denominators in AdamW second moments, and attention branches whose forward and backward paths can diverge if quantized inconsistently.

The main claim is that these failures are module-specific enough to need module-specific repairs. That is a useful framing. It turns “4-bit pretraining” from a single compression setting into a set of precision contracts: where to quantize, where to transform values before storage, and where BF16 remains necessary.

Technical Approach

For linear projections, the method adds a lightweight LoRA-SVD decomposition to reduce direct quantization error. The paper reports that this breaks a direct-quantization ceiling: the loss gap for linear-only quantization falls from 1.40% to 0.61%. The point is not that low-rank adapters replace the base projection, but that they absorb error that NVFP4 alone cannot represent cleanly.

For optimizers, Full-Stack FP4 treats AdamW and Root (Muon) differently. AdamW’s second moment values are non-negative and heavy-tailed, so directly storing them in NVFP4 produces unstable denominator behavior. The authors instead transform the second moment before storage. For Root, they support native NVFP4 Newton-Schulz iterations, which extends the quantized stack into optimizer arithmetic rather than only optimizer storage.

Attention receives the most conservative treatment. The proposed scheme quantizes Q, K, V, and backward dS, while keeping sensitive paths in BF16. The paper identifies PV and dOVᵀ branches as especially error-prone, and pairs the mixed-precision rule with unified tensor reuse so forward and backward computations remain aligned. That alignment matters because attention has coupled computation paths; a small mismatch in one branch can be amplified by the next.

Results and Analysis

The central experiment is a 3B-parameter model trained on 64B tokens, with 24 layers, hidden size 3072, FFN dimension 9216, 24 attention heads, 12 KV heads, and a 2048-token context window. The training recipe uses Root+AdamW, peak learning rate 1.5 × 10⁻⁴, a 3.0 × 10⁻⁴ peak learning rate for the LoRA branch, 3B warmup tokens, 1M-token batches, gradient clipping at 1.0, and cosine decay.

On that setup, Full-Stack FP4 ends with a 1.47% loss gap relative to the BF16 Root+AdamW baseline. Figure 5 shows the training-loss comparison: the FP4 run tracks the BF16 baseline closely rather than showing the late instability expected from naïve full-stack quantization.

The ablations support the paper’s main diagnosis. Linear-only NVFP4 improves when LoRA-SVD is added, cutting the reported loss gap from 1.40% to 0.61%. In the appendix loss-surface study, the high-precision baseline keeps a smoother optimization path, the proposed quantized stack becomes steeper but remains trainable, and the naïve NVFP4 baseline collapses by step 672. That visualization is not a replacement for downstream evaluation, but it is consistent with the numerical story: unguarded 4-bit training does not merely degrade slowly; it can enter invalid training states.

The evidence is persuasive for the narrow claim that a carefully partitioned FP4 stack can pretrain a 3B model for 64B tokens with small loss degradation. It is less conclusive for deployment-wide claims. The paper reports pretraining loss rather than downstream benchmark quality, evaluates one main model scale, and does not show whether the same precision rules hold for longer context, larger models, or different data mixtures. The result is still practically relevant: teams trying to reduce pretraining memory and arithmetic cost get a concrete map of which transformer paths can be quantized and which should remain protected.