TetraJet-v2: Accurate NVFP4 Training for Large Language Models with Oscillation Suppression and Outlier Control

AI summary

Overview

Research area: Efficient machine learning — low-precision (4-bit) fully quantized training of large language models.

Technical level: Advanced. The paper assumes familiarity with floating-point formats, quantization, and LLM pre-training dynamics, though the core ideas are describable in plain terms.

Scope: TetraJet-v2 is an end-to-end NVFP4 training method that targets the two optimization failures specific to 4-bit LLM training — weight oscillation and activation outliers — and closes much of the remaining accuracy gap to BF16 while running faster than FP8.

What This Paper Is About

Training LLMs is extremely expensive, and pushing arithmetic down from 16-bit or 8-bit to 4-bit formats like NVFP4 promises large speed and memory wins. But prior 4-bit training either degrades badly, silently keeps some layers in higher precision, or only works in mixed settings. This paper asks why fully 4-bit training fails and builds algorithmic fixes for the two specific culprits it identifies.

Key Contributions

1. An unbiased double-block NVFP4 linear layer. A 1×128 outer block with a global scale plus 1×16 inner blocks, deterministic rounding in the forward pass and stochastic rounding in the backward pass, and gradient estimation aligned to the quantized activations and weights rather than the master values. This yields unbiased gradient estimates and better convergence than NVIDIA's NVFP4 recipe.
2. OsciReset, the first effective weight-oscillation suppressor for LLM pre-training. Oscillating master weights (identified via an oscillation-risk statistic) are reset to the center of the FP4 bin they currently occupy, which leaves the quantized weights unchanged at that instant but lets them re-enter optimization on a trajectory less prone to flipping.
3. OutControl, a mix-precision outlier-handling scheme. Random Hadamard Transformation (RHT) is applied only in the backward pass (forward RHT hurt optimization in their experiments), and a small, statically selected set of activation outlier channels is retained in FP8 across forward and backward.
4. Efficient CUDA kernels and a rigorous fully-4-bit evaluation. All linear layers in all Transformer blocks are quantized for both their method and the baselines, confirming 1.67× end-to-end speedup over FP8 with only minimal kernel overhead.

Main Findings

• Both failure modes are optimization problems, not just quantization error. The authors argue the gap between FP4 and high-precision training cannot be explained by representational error alone; weight oscillation and outlier sensitivity are distinct optimization pathologies.
• Weight oscillation is a late-training phenomenon. As the learning rate approaches zero, an increasing fraction of latent weights (w/s) sit near quantization decision thresholds, so tiny gradient steps flip the quantized value between adjacent FP4 bins while the master weight barely moves.
• OsciReset beats prior oscillation remedies. Q-EMA smoothing, freezing oscillating weights, and adding a threshold-repulsion penalty (Dampen) — all tuned for Vision Transformers — each severely degraded final LLM performance. Only OsciReset consistently improved it.
• RHT helps in the backward pass but hurts in the forward pass. This contradicts both Castro et al. (who rotate the forward) and NVIDIA (who apply RHT only to dW); the authors found RHT improves both dX and dW computations.
• Activation quantization in the forward pass is the single largest accuracy bottleneck. In loss decomposition on OLMo2-370M (52B tokens), quantizing only the forward activation cost +0.44 PPL, versus +0.22 for forward weights; quantizing dX and dW individually was nearly free (+0.07 and −0.01) but together cost +0.21, suggesting error accumulation.
• MLP layers are more quantization-sensitive than attention layers. Leaving mlp.ffn2 in BF16 recovered −0.35 PPL, more than leaving out qkv (−0.15) or attn.out (−0.16).
• Substantial accuracy recovery. On OLMo2-370M trained for 212B tokens, TetraJet-v2-full reached 43.60 average downstream accuracy versus 44.13 for BF16, 42.79 for Quartet, and 42.77 for the NVIDIA recipe — narrowing the BF16 gap by 51.3% on average across model sizes.
• Speed. CUDA kernels on an RTX 5090 delivered 1.94× (base) and 1.83× (full) speedups on linear layers versus TransformerEngine FP8, translating to 1.75× and 1.67× end-to-end at the Transformer-layer level.

Methodology in Plain English

The authors start from the standard quantized-training setup: keep a high-precision "master" copy of the weights, quantize activations, weights, and gradients for every matrix multiply, and update the master weights with the resulting gradients.

They first redesign the linear layer itself. NVFP4 pairs 4-bit values with an 8-bit scale factor, but that scale factor can only represent a limited range, so they wrap a coarse 1×128 block around the fine 1×16 groups. They also make sure the backward pass is a statistically unbiased estimate of the true gradient — using stochastic rounding on gradients, and computing weight gradients against the already-quantized activations, so forward and backward see the same quantities. This is a small change with real theoretical consequences: unbiased gradients are what allow SGD to converge.

For oscillation, they track how far each master weight and its quantized counterpart travel over a window. A high ratio of quantized distance to master distance means the weight is thrashing between bins. Tracking every weight is expensive, so they track only the 5% closest to thresholds — the weights that actually oscillate — costing about 0.6 bytes per parameter. When an element's oscillation risk crosses a threshold, they snap the master weight to the center of its current bin. Crucially, the quantized weight value does not change at that moment, so there is no immediate accuracy hit; the weight simply gets a fresh starting point and can be optimized toward a more stable value instead of ping-ponging forever.

For outliers, they rotate tensors with a randomized Hadamard transform before quantizing, but only for the two backward matrix multiplies — empirically, rotating the forward pass added more error than it removed. They also observe that activation outliers are structural: the same handful of channels are large across different inputs and training steps. That means they can run a short calibration at 1% into training, pick the top 10% of channels by accumulated norm, and hold those channels in FP8 permanently for both forward and backward. This static selection avoids per-batch dynamic decisions and keeps training and inference aligned.

Finally, they validate everything by pre-training OLMo-2 models at 70M, 150M, and 370M parameters on 52B, 107B, and 212B tokens, comparing against BF16, Quartet (MXFP4), and NVIDIA's NVFP4 recipe — with all methods forced into a strictly fully-4-bit setting for fairness.

Why This Matters

The paper reframes 4-bit training as an optimization problem rather than a pure compression problem, which opens a different line of attack than simply shrinking quantization error. It also demonstrates that a fully 4-bit recipe can beat a competitive FP8 baseline in wall-clock terms while staying close to BF16 accuracy, which is a meaningful data point for the feasibility of FP4 pre-training.

Real-world applications:

• Cost reduction for LLM pre-training and continued pre-training, where compute budgets routinely run into tens of millions of dollars per model.
• Memory-constrained fine-tuning and domain adaptation on smaller GPU fleets, where 4-bit weights and gradients roughly halve memory versus FP8.
• On-device or edge inference models that are trained and shipped in the same low-precision regime, avoiding a separate quantization step.
• Distributed training at scale, since communication volume for activations and gradients shrinks proportionally with precision.

Industry relevance: The method targets NVIDIA Blackwell's native NVFP4 support, is released as open code, and was developed in a Tsinghua–Bosch collaboration — a combination that suggests near-term applicability in commercial training stacks. The 1.67× over FP8 figure is the kind of number that directly translates into hardware-hour savings at scale.

Future Directions

• Scaling beyond 370M parameters. All experiments are on models up to 370M parameters; whether OsciReset and OutControl retain their effectiveness at multi-billion-parameter scale and longer token budgets is untested.
• Extending the loss-decomposition analysis. The finding that forward activation quantization dominates the error, and that MLP layers are the most sensitive modules, hints that selective precision allocation across layers could yield further gains — but the paper only begins this investigation.
• Reconciling the RHT placement question. The authors' finding that backward RHT helps both dX and dW while forward RHT hurts directly contradicts prior work; resolving why could lead to better rotation strategies.
• Better oscillation detectors. The current approach relies on a hand-chosen risk threshold and a calibration window; whether oscillation risk can be predicted or adaptively thresholded remains open.

Target Audience

Researchers and engineers working on low-precision training systems, quantization for LLMs, and efficient ML infrastructure will get the most from this paper. Practitioners running large-scale pre-training on Blackwell-class hardware will find the concrete recipe and kernel results directly useful. Readers interested in optimization dynamics under quantization — particularly the weight oscillation analysis — will find the diagnostic framing valuable even outside the FP4 setting. A background in quantization and Transformer training is assumed.

Authors’ abstract

Large Language Models (LLMs) training is prohibitively expensive, driving interest in low-precision fully-quantized training (FQT). While novel 4-bit formats like NVFP4 offer substantial efficiency gains, achieving near-lossless training at such low precision remains challenging. We introduce TetraJet-v2, an end-to-end 4-bit FQT method that leverages NVFP4 for activations, weights, and gradients in all linear layers. We identify two critical issues hindering low-precision LLM training: weight oscillation and outliers. To address these, we propose: 1) an unbiased double-block quantization method for NVFP4 linear layers with practically optimal convergence in LLM training, 2) OsciReset, the first effective algorithm to suppress LLMs' weight oscillation bottleneck, and 3) OutControl, a mix-precision algorithm to retain outlier accuracy. TetraJet-v2 outperforms prior methods on FP4 pre-training for LLMs across models up to 370M parameters trained up to 212B tokens, reducing the performance gap to BF16 by an average of 51.3% while enabling an 1.67x end-to-end speedup over FP8. The code is available at https://github.com/thu-ml/TetraJet-v2-NVFP4Training.