🤖 machine learning

Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention

This paper introduces Full-Stack FP4, a modular framework that extends NVFP4 quantization beyond linear projections to include stable optimizer states, Root/Muon computations, and mixed-precision attention, achieving near-BF16 pretraining performance with significant memory and speedup gains on a single RTX 5090.

Original authors: Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

Published 2026-08-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, digital brain how to speak, write, and think. This "brain" is a Large Language Model (LLM), and to learn, it has to process mountains of information. But there's a catch: the more it learns, the more hungry it gets for computer memory and electricity. It's like trying to run a massive, high-definition movie on a tiny, old laptop; the screen gets blurry, and the fan screams. To solve this, scientists have been trying to shrink the numbers the computer uses to do its math. Instead of using super-precise, heavy numbers (like a detailed photograph), they've been switching to smaller, lighter numbers (like a pixelated sketch). This is called "low-precision training." Recently, a new type of tiny number called "FP4" (4-bit floating point) has arrived, promising to make these giant brains much smaller and faster. But here's the problem: while the main parts of the brain can handle these tiny numbers, the "helper" parts—like the memory it keeps for learning and the complex math it does to understand sentences—keep crashing or getting confused when forced to use such small numbers. It's like having a race car engine but trying to run it on a bicycle chain; the engine is ready, but the rest of the machine can't keep up.

This paper introduces a clever new toolkit called Full-Stack FP4, which acts like a master mechanic for these digital brains. The researchers realized that you can't just use the same tiny-number rules for every single part of the brain. Some parts, like the main connections, can handle the tiny numbers just fine. But other parts, like the "memory" the computer uses to remember what it learned yesterday, are very sensitive and need a special kind of care. The team created a modular system where they use different "recipes" for different parts of the brain. For the main connections, they use a trick called LoRA-SVD, which is like keeping a high-resolution sketch of the most important details while using pixelated blocks for the rest. For the memory, they invented a way to "smooth out" the numbers before squishing them into tiny boxes so they don't get lost. For the math that helps the brain focus on the right words (attention), they decided to keep the most sensitive parts in high-definition while letting the rest stay pixelated.

When they tested this full system on a 3-billion-parameter model (a medium-sized digital brain) using 64 billion words of training data, the results were impressive. The "Full-Stack FP4" brain learned almost exactly as well as the traditional, heavy-duty version. The difference in their learning scores was a tiny 0.838%, which is barely noticeable. In fact, when they tested how well the brain could answer questions it had never seen before, the new tiny-number version was actually slightly better at guessing the right words, even though it was using much less memory. On a single powerful graphics card (an RTX 5090), this new method made the brain's learning process 2.5 to 2.8 times faster for certain tasks and used 38% to 42% less memory than the standard high-precision method.

The researchers are careful to note that while this works great for models up to 3 billion parameters, they haven't proven yet if it will work for even bigger brains or over much longer training times. They also found that you can't just throw tiny numbers at everything; if you try to force the most sensitive parts of the brain to use the smallest numbers without these special recipes, the learning breaks down. But for now, this "Full-Stack" approach shows that we can build smarter, faster, and cheaper AI by treating different parts of the brain with the specific level of precision they actually need, rather than forcing the whole system to be one size fits all.

Technical Summary: Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention

1. Problem Statement

Recent advancements in NVIDIA Blackwell GPUs have made native NVFP4 (4-bit) Tensor Core support a practical target for Large Language Model (LLM) pretraining. While existing research has successfully demonstrated stable 4-bit pretraining for Transformer linear projections, significant gaps remain in other critical training modules. Most current NVFP4 recipes retain higher-precision (e.g., BF16 or FP8) for optimizer states, optimizer computations, and attention mechanisms.

The primary challenge is that these modules possess distinct numerical structures that a uniform quantization rule cannot address:

Linear Projections: Gradient signals weaken over training while quantization noise remains constant, reducing the efficacy of continued low-precision training.
AdamW States: The second-moment momentum is non-negative, heavy-tailed, persistent, and sensitive to denominator values, making direct 4-bit storage prone to distortion.
Root/Muon Optimizers: These rely on repeated Newton–Schulz matrix products where low-precision iteration errors can accumulate and destabilize convergence.
Attention: The operator contains softmax-sensitive paths (e.g., $PV$, $dOV^\top$ ) that require high precision, while other paths (e.g., $QK^\top$ ) are more robust to quantization.

Consequently, there is a lack of a unified, modular framework that enables stable, native NVFP4 pretraining across the entire training stack (projections, optimizers, and attention) without relying on mixed-precision fallbacks for critical components.

2. Methodology: Full-Stack FP4

The authors present Full-Stack FP4, a modular framework comprising four independent, composable recipes designed to address the specific numerical difficulties of each training module.

2.1 LoRA-SVD Linear Projections

To mitigate the degradation of gradient signals in later training stages, the authors propose LoRA-SVD.

Mechanism: This method parameterizes weights as $W = W_{res} + \beta L_2 L_1$ . It maintains a compact principal subspace ( $L_2 L_1$ ) in BF16 while computing the dominant residual ( $W_{res}$ ) in NVFP4.
Backward Pass: A subspace-constrained backward pass uses a Cholesky-QR constraint to ensure the BF16 low-rank branch updates align with the principal subspace, while the residual updates the NVFP4 component.
Realignment: The subspace is periodically realigned (every 2,048 steps) using randomized SVD to reset the components, ensuring the BF16 subspace remains optimal.

2.2 Quantized-State AdamW (Q-AdamW)

To stabilize the storage of persistent, heavy-tailed momentum states in NVFP4:

First Moment ( $m_t$ ): Applies Hadamard mixing to smooth the distribution before Stochastic Rounding (SR) and adaptive 4/6 quantization.
Second Moment ( $v_t$ ): Utilizes a three-stage preprocessing pipeline before NVFP4 storage:
1. Square Root: Shortens the heavy tail.
2. Tile Mean: Separates the positive common component (stored in BF16) from an approximately zero-mean residual.
3. Hadamard Transform: Spreads localized energy in the residual to make it more symmetric.
Reconstruction: The state is reconstructed by inverting the transform and mean separation, then squaring the result. This pipeline avoids the need for SVD-based subspace operations.

2.3 NVFP4 Root Optimizer (Q-Root)

For the Newton–Schulz (NS) iterations used in the Root optimizer:

Direct NVFP4 Execution: Instead of a mixed-precision route, the NS matrix products are executed directly in NVFP4.
Stabilization: Stability is achieved through shape-dependent coefficients (optimized $a, b, c$ for specific matrix aspect ratios) and p99 outlier clipping. This approach limits error amplification across successive NS steps without requiring auxiliary mixed-precision paths.

2.4 Mixed-Precision Attention (Q-Attn)

To balance precision and efficiency in attention:

Precision Split: The framework assigns NVFP4 to robust paths ( $QK^\top$ and $dS$) and retains BF16 for softmax-sensitive paths ($PV$, $P^\top dO$ , and $dOV^\top$ ).
Consistency: Quantized views of $Q, K, V$ are reused between forward and backward passes to ensure tensor consistency and avoid redundant memory traffic.

3. Key Contributions

LoRA-SVD Projections: A mergeable reparameterization that reduces the linear-only loss gap from 1.40% to 0.61% compared to standard NVFP4, protecting a compact subspace in BF16.
Quantized-State AdamW: The first demonstration of stable NVFP4 AdamW momentum-state storage in an integrated pretraining run, utilizing a transformed-state pipeline that reduces second-moment reconstruction error by ~60%.
NVFP4 Root: A direct low-precision Newton–Schulz path enabled by shape-aware coefficients and clipping, reducing iteration error from ~10% to ~3%.
Trainable Mixed-Precision Attention: A strategy that protects softmax-sensitive paths in BF16 while quantizing suitable forward and backward paths, enabling a consistent NVFP4 pretraining operator.

4. Experimental Results

The authors evaluated the framework on a 3B parameter Transformer trained for 64B tokens.

Training Loss: The BF16 baseline reached a loss of 2.267, while Full-Stack FP4 reached 2.286, resulting in a relative gap of 0.838%.
Downstream Performance: Zero-shot evaluation showed Full-Stack FP4 had an average perplexity of 26.665 (vs. 26.675 for BF16) and an average accuracy 0.10 percentage points lower than the BF16 baseline.
Ablation Studies:
- LoRA-SVD alone reduced the linear-only loss gap from 1.40% to 0.61%.
- The modular recipes can be enabled independently or combined, with "Full-Stack" representing the aggregate coverage.
Native Hardware Efficiency (RTX 5090):
- Root Speedup: Full-Stack FP4 achieved 2.50×–2.83× speedups over optimized BF16 and 1.83×–2.17× over NVIDIA Transformer Engine (TE) for the Root optimizer step.
- Memory Savings: The framework reduced AdamW peak memory by 37.9%–42.5% compared to BF16.
- Time Savings: Estimated matched-loss time savings were calculated at 56.9%–61.9% when accounting for the reduced step time and the slight increase in tokens required to match BF16 loss.

5. Significance and Limitations

Significance:
The paper demonstrates that a "Full-Stack" approach to NVFP4 pretraining is feasible, covering projections, optimizer states, optimizer computation, and attention. It establishes that native 4-bit pretraining can achieve performance comparable to BF16 (within <1% loss gap) while offering substantial memory and throughput benefits on Blackwell-class hardware. The work provides a modular recipe set that allows practitioners to adopt low precision incrementally.

Limitations:

Scale: Experiments are limited to models up to 3B parameters and 64B tokens. The paper does not establish behavior at 7B+ scales or over longer training horizons.
Efficiency Measurement: Native hardware efficiency was measured on four decoder blocks on a single RTX 5090, not on a complete distributed model.
Estimates: Time savings calculations are cross-setting estimates based on fixed batch sizes and do not account for potential throughput gains from increasing batch size using freed memory.
Kernel Fusion: Some kernels (LoRA backward, attention) are not yet fully fused, and the implementation does not yet utilize all FlashAttention-4 optimizations.
Format Specificity: The recipes and measurements are specific to NVFP4 and do not claim transferability to MXFP4.