PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuxiang Chen, Yifan Liu, Xiaoming Xu, Pengle Zhang, Michael Beyer, Martin Rapp, Jun Zhu, Jianfei Chen

arXiv:2510.27527v3cs.LGcs.AI

TL;DR

Fully quantized 4-bit LLM training is attractive for reducing training cost but remains difficult because NVFP4 introduces weight oscillation and outlier-related accuracy problems. TetraJet-v2 combines unbiased NVFP4 linear layers, OsciReset, and OutControl, reducing the BF16 performance gap by 51.3% while achieving 1.67× speedup over FP8 in its evaluated setting.

  • Problem

    Fully quantized NVFP4 training must overcome weight oscillation and outlier sensitivity to approach high-precision LLM training.

  • Method

    TetraJet-v2 uses NVFP4 for activations, weights, and gradients, with unbiased double-block quantization, OsciReset, and mixed-precision OutControl.

  • Results

    51.3%: TetraJet-v2 narrows the BF16 performance gap and achieves 1.67× end-to-end speedup over FP8 on models up to 370M parameters trained for up to 212B tokens.

  • Takeaways & Limitations

    The method enables stable, accurate fully NVFP4 LLM training while retaining a practical speed advantage over FP8 in the evaluated setting.

  • Takeaways & Limitations

    Experiments are limited to OLMo-2 models up to 370M parameters and data scales up to 212B tokens, leaving larger models, datasets, and architectures for future work.

Abstract

from arXiv · show

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.

1. Introduction

TetraJet-v2 targets the instability and accuracy barriers that make fully quantized NVFP4 training difficult, combining unbiased linear-layer quantization with oscillation suppression and outlier control. On OLMo-2 models up to 370M parameters and 212B tokens, it narrows the BF16 performance gap by 51.3% and achieves 1.67× end-to-end speedup over FP8.

  • Challenges: Weight oscillation and outlier sensitivity are identified as major obstacles beyond quantization error alone, with oscillation disrupting optimization through repeated bin changes.Small updates to high-precision master weights can repeatedly flip their quantized representations between adjacent bins.
  • Method: TetraJet-v2 applies NVFP4 to activations, weights, and gradients in all linear layers while targeting the optimization failures specific to fully FP4 LLM training.The method is designed as an end-to-end fully quantized training approach rather than a partially high-precision recipe.
  • Contributions: The proposed linear-layer design uses unbiased double-block NVFP4 quantization and an empirically optimal fully-FP4 configuration with Random Hadamard Transform.These choices are intended to improve convergence in the LLM training setting.
  • Contributions: OsciReset suppresses LLM weight oscillation, while OutControl retains outlier accuracy through mixed precision in forward and backward computation.Together, the methods address the two principal optimization and representation bottlenecks identified by the paper.
  • Evaluation: 51.3%: TetraJet-v2 narrows the performance gap to full precision on OLMo-2 models up to 370M parameters trained for 212B tokens, while achieving 1.67× end-to-end speedup over FP8.The evaluation covers fully NVFP4 pre-training with all proposed methods integrated.

2. TetraJet-v2’s NVFP4 Linear Layer Design

The NVFP4 linear layer uses nested scaling, deterministic forward rounding, stochastic backward rounding, and aligned tensor layouts across the forward and backward matrix multiplications. This design supports hardware-efficient FP4 computation while targeting unbiased gradient estimation and improved accuracy over NVIDIA’s recipe.

  • Rounding: Deterministic rounding is used in the forward pass to minimize quantization error, whereas stochastic rounding is used in backward computation to obtain unbiased gradient estimates.This distinction is central to the method’s convergence-oriented quantization design.
  • Double-block quantization: NVFP4 represents values with 1 × 16 inner blocks and an additional 1 × 128 outer block, using finer scaling to improve accuracy and hardware friendliness.The outer block avoids full-tensor maximum computation and performs better than NVIDIA et al. (2025)'s per-tensor second scaling in the reported experiments.
  • Linear-layer computation: The linear layer computes Y = X × W⊤, dX = dY × W, and dW = dY⊤× X, quantizing the six inputs to these three matrix multiplications in NVFP4.The chosen group shapes are 1 × 16 for some operands and 16 × 1 for the transposed-orientation operands to meet acceleration requirements.
  • Design differences: The design aligns dW computation with the quantized forward input bX, unlike NVIDIA et al. (2025), which estimates X in backward computation.The alignment is paired with stochastic backward rounding and a 1 × 16 weight-block shape.
  • Design outcome: The resulting linear layer is reported to ensure unbiased gradient estimation and to outperform previous NVFP4 training designs.The analysis connects unbiasedness to convergence of SGD and the experiments report better NVFP4 training performance.

3. OsciReset: Oscillation Suppression Through Resetting Master Weights

OsciReset addresses weight oscillation in NVFP4 training by detecting weights that repeatedly cross quantization thresholds and resetting them to bin centers. The method targets late-training instability while preserving optimization opportunities and adding little overhead.

  • 3.1. Oscillation Phenomenon in LLMs Pre-training: Weight oscillation arises when small master-weight updates repeatedly move quantized weights across adjacent bins, disrupting convergence near the end of training.The proportion of latent weights near quantization thresholds grows during training, explaining increasingly frequent quantized-value switches.
  • 3.2. Detecting Oscillating Weights: OsciReset identifies oscillating elements using the ratio of quantized-weight movement to master-weight movement, with larger ratios indicating higher oscillation risk.The implementation can track only the top-5% of weights closest to thresholds, adding approximately 0.6 Bytes/param.
  • 3.3. Reducing Oscillation by Resetting Master Weights: The method stabilizes detected weights by resetting each master weight to the center of its current quantization bin when the learning rate is low.Resetting leaves the immediately quantized weights unchanged, avoiding an immediate performance drop while allowing future optimization.
  • 3.3. Reducing Oscillation by Resetting Master Weights: Resetting is preferable to freezing because oscillating parameters can pursue new optimization trajectories instead of remaining trapped between two quantized bins.The resulting trajectories are less prone to oscillation and are reported to improve model performance.
  • 3.3. Reducing Oscillation by Resetting Master Weights: Unlike prior oscillation methods aimed mainly at Vision Transformers, OsciReset applies directly to LLM pre-training from scratch with minimal tuning and negligible latency and memory overhead.The section emphasizes that suppressing oscillation must preserve optimization across all parameters rather than simply immobilizing unstable elements.

4. OutControl: Outlier Control for Activation and Gradients

OutControl addresses activation and gradient outliers in NVFP4 training by combining backward-pass Random Hadamard Transformations with higher-precision retention for persistent activation channels. Static channel selection exploits the structural consistency of outliers across inputs and training steps.

  • 4. OutControl: Outlier Control for Activation and Gradients: Activation and gradient outliers inflate quantization scales, pushing small values toward zero and reducing their representational accuracy in FP4.This outlier sensitivity is a central obstacle to accurate low-precision training.
  • 4.1. Random Hadamard Transformation for Quantized Backpropagation: Random Hadamard Transformation is applied to both operands of the backward dX and dW matrix multiplications, preserving mathematical equivalence while controlling quantization variance.The forward pass omits Hadamard transformation because the authors observe that it harms optimization.
  • 4.2. Precision Retaining for Activation Outliers: OutControl statically retains persistent high-variance activation channels in higher precision while quantizing the remaining channels, improving forward accuracy and combining with RHT for backward control.The selected outlier pattern is structural: a small set of channels consistently has larger norms across inputs and training steps.
  • 4.2. Precision Retaining for Activation Outliers: Outlier channels are selected offline by accumulating activation ℓ2-norms during an early calibration window and choosing the top p% channels, typically p = 10% for FP8.Fixed selection removes per-batch or per-step dynamic adjustment and aligns the training structure with inference.
  • 4.2. Precision Retaining for Activation Outliers: The mixed-precision formulation allows retained activation channels and corresponding weights to use FP8, BF16, or FP4, while the complement remains quantized.The design is described as producing unbiased gradient estimates, supporting SGD convergence and improved NVFP4-training performance.

5. Experiments

Experiments evaluate fully quantized NVFP4 training across OLMo2 model sizes, compare FP4 methods, and analyze accuracy, oscillation suppression, efficiency, and precision switching. TetraJet-v2-full achieves the strongest reported FP4 results while improving speed and identifying activation quantization as a key bottleneck.

  • Results: TetraJet-v2-full reaches the lowest PPL and the highest average downstream performance among FP4 methods, further closing the gap to high-precision training.The comparison covers OLMo2-370M trained for 212B tokens and reports training and validation PPL plus downstream benchmarks.
  • NVFP4 Configuration: NVFP4 linear layers surpass NVIDIA’s design through refined scaling and unbiased gradient estimation, supported by ablations over block size, aligned activations, stochastic rounding, and Hadamard transforms.The experiments use strict full quantization of activations, weights, and gradients in all Transformer linear layers.
  • Method Combination: OutControl improves accuracy throughout training, while OsciReset adds improvement when oscillation suppression begins at approximately 150B tokens.The combination is evaluated with the OLMo2-370M validation-loss curve trained for 212B tokens.
  • Oscillation Suppression: Among oscillation-suppression methods, only OsciReset consistently maintains a substantial optimization improvement for LLMs; Q-EMA, Freeze, and Dampen hurt final performance.All methods begin suppression at approximately 65B tokens in the OLMo2-150M, 107B-token comparison.
  • Efficiency: TetraJet-v2 accelerates Transformer-layer linear computations by 1.94× in base and 1.83× in full configurations, yielding 1.75× and 1.67× end-to-end speedups over FP8.Measurements use an RTX 5090 with MicroBatchSize=4, SeqLen=1 024, and Hidden=16 384.
  • Loss Decomposition: Activation quantization is the main forward accuracy bottleneck, while jointly quantizing dX and dW accumulates more error than quantizing either gradient computation alone.The decomposition also finds higher sensitivity in MLP linear modules, especially mlp.ffn2, than in attention modules.

6. Discussion, Conclusion and Limitations

TetraJet-v2 extends fully NVFP4 LLM training with methods targeting oscillation, unbiased quantization, and outlier sensitivity. The discussion reports practical efficiency and identifies scope boundaries for current validation.

  • 6.1. Discussion: OsciReset addresses an LLM-specific optimization failure that prior vision-transformer oscillation methods do not transfer to effectively.The paper attributes this difference to contrasting data and optimization behavior, including much longer token-per-parameter training horizons for LLMs.
  • 6.1. Discussion: The double-block scaling design uses a 1 × 128 outer block with FP32 scaling to avoid a full extra scan, but remains an implementation detail.Larger future models may require larger granularity to balance scale overhead and runtime efficiency.
  • 6.1. Discussion: RHT can degrade forward NVFP4 quality yet improve dW and dX computation when backward quantization uses unbiased gradients.Its effectiveness remains under exploration for MoE architectures and larger-scale training.
  • 6.2. Conclusion: TetraJet-v2 combines unbiased double-block quantization, oscillation reduction, and outlier precision control to enable stable and accurate fully NVFP4 LLM training.The method uses NVFP4 for activations, weights, and gradients in all linear layers.
  • 6.3. Limitations: Experiments cover OLMo-2 models up to 370M parameters and data scales up to 212B tokens because of limited compute resources.The authors identify larger models, more data, and additional architectures as future directions.

Impact Statement

Low-precision training may reduce the hardware cost and energy use of large-scale pre-training, while also lowering barriers to developing and deploying powerful language models. These benefits coexist with potential misuse risks.

  • Impact Statement: Low-precision training may reduce hardware cost, energy consumption, and the carbon footprint of large-scale pre-training.The statement frames these effects as potential benefits of improved training efficiency.
  • Impact Statement: More efficient training may lower barriers to developing and deploying powerful language models, potentially facilitating misinformation, spam, or misuse at scale.The passage presents these as possible harmful uses rather than established outcomes.
  • Low-Precision Training: Fully quantized training differs from post-training quantization and quantization-aware training by quantizing activations, weights, and gradients during both forward and backward passes.This training-oriented design targets computational efficiency rather than only post-training inference efficiency.
  • Outlier Control: Fine-grained block-wise quantization reduces the impact of outlier values by confining their effect to smaller groups, but outlier handling remains an algorithmic challenge.Architectural changes such as gated attention or attention sinks do not fully mitigate outliers.

B.2. Hyperparameters for NVFP4 Training Techniques

The NVFP4 training techniques section specifies implementation parameters and the periodic workflow used to track and suppress oscillation. The procedure quantizes weights, identifies sensitive elements, accumulates statistics, and periodically applies suppression.

  • B.2. Hyperparameters for NVFP4 Training Techniques: Oscillation statistics compare FP4 quantized-dequantized weights with master-weight snapshots and aggregate element-wise distances into tensor-level measures.The tracked quantities include master and quantized snapshots for selected weights.
  • B.2. Hyperparameters for NVFP4 Training Techniques: The memory-efficient variant samples sensitive indices by ranking weights with s_k = |W_k − Q(W_k)|/binwidth(Q(W_k)).It selects the top-psample% elements and tracks their accumulated statistics.
  • B.2. Hyperparameters for NVFP4 Training Techniques: The training loop computes loss and gradients, updates parameters with an optimizer step, then conditionally tracks statistics and calls OscillationSuppress.The suppression routine is implemented as a separate algorithm invoked at the configured period boundary.
  • B.2. Hyperparameters for NVFP4 Training Techniques: OsciReset periodically computes oscillation statistics after optimizer updates and invokes suppression when the configured period reaches its accumulation boundary.The trainer uses Tstart, Tperiod, Taccu, and τosci to control when tracking and suppression occur.

C.2. Overhead & Memory Efficiency Analysis

OsciReset has low measured runtime and memory overhead, while near-threshold tracking preserves its statistical behavior. The analysis also shows why high quantization error alone is insufficient for selecting weights to reset.

  • C.2. Overhead & Memory Efficiency Analysis: Only ∼0.4s per 200 steps, or less than 0.06% of total training time, is added by periodic resetting on 150M-OLMo2 training.The measurement uses 8× RTX5090 hardware and Tperiod = 200 steps.
  • C.2. Overhead & Memory Efficiency Analysis: Tracking only the top-5% weights nearest quantization thresholds adds ∼0.6 Bytes per parameter and has negligible impact on model quality.The approximation preserves oscillation statistics and can be naturally sharded across GPUs without additional communication cost.
  • C.3. Discussion on Oscillation-Quantization Error Relationship: ∼70% of harmful oscillators lie far from quantization-bin centers, supporting resets that remove high quantization error and alter their optimization trajectory.The analyzed harmful oscillators comprise fewer than 5% of weights.
  • C.3. Discussion on Oscillation-Quantization Error Relationship: Only ∼15% of high-quantization-error weights oscillate, so resetting every such weight would disrupt normal optimization.Explicitly identifying oscillating weights is therefore necessary to target the harmful subset.

D.1. NVFP4 Linear Layer Design

The NVFP4 linear layer combines finer-grained scaling, unbiased gradient estimation, and selective Hadamard transforms, while OsciReset and OutControl address oscillation and outliers during training.

  • NVFP4 Linear Layer Design: Random Hadamard Transformation improves both dX and dW computations in the unbiased linear framework, whereas forward Hadamard Transformation is harmful.This differs from the NVIDIA recipe, which reported benefits mainly for dW.
  • NVFP4 Linear Layer Design: TetraJet-v2’s NVFP4 linear design surpasses NVIDIA’s design through refined scaling and unbiased gradient estimation.The ablation attributes the improvement to finer outer scaling, forward/backward alignment of bX, and stochastic rounding.
  • Oscillation Suppression: OsciReset causes a substantial decay in oscillating weights, countering the drastic increase observed without suppression.Oscillating weights are identified using OsciRisk(w), with the cited figure using a threshold of 16.
  • Oscillation Suppression: OsciReset is generally insensitive to its oscillation threshold and suppression start point across model sizes and data scales.The experiments use τosci = 8 and Tstart ≈60% · Tmax consistently from 70M to 370M models.
  • Outlier Control: OutControl’s static channel selection improves both forward and backward passes compared with random precision-retention choices.This supports fully NVFP4 training with quantization applied in both directions.

E. More FP4 Baselines Results

Expanded comparisons show that TetraJet-v2 consistently outperforms recent FP4 training methods, including under the Muon optimizer.

  • Expanded Baseline Comparison: TetraJet-v2 consistently outperforms all three additional recent state-of-the-art baselines under fully-FP4 settings where applicable.The comparison uses validation perplexity, for which lower is better.
  • Muon Optimizer: With Muon, TetraJet-v2 consistently improves over the NVIDIA baseline, and TetraJet-v2-full achieves the best validation perplexity among the compared 4-bit methods.The evaluation covers OLMo-2-70M and OLMo-2-150M trained with 52B tokens.

G. Detailed Efficiency Results for TetraJet-v2

TetraJet-v2’s finer-grained scaling adds little overhead, while its kernels preserve throughput and end-to-end speedups over FP8 at Transformer-block and full-model scales.

  • Latency Overhead: The 1×16 weight-scaling scheme adds only 2.49% average overhead compared with the 16×16 baseline across evaluated hidden sizes and sequence lengths.The comparison changes only the weight quantization recipe, although the baseline benefits from avoiding backward weight re-quantization.
  • Linear-Layer Throughput: TetraJet-v2-full has slightly lower single-layer throughput than the base version because of its mix-precision design, despite accounting for full quantization and transform overhead.The throughput benchmarks include quantization, de-quantization, and Random Hadamard Transform costs.
  • Transformer-Block Speedup: The complete Transformer-block benchmarks report end-to-end speedup ratios against TE-FP8 across different microbatch sizes and sequence lengths.The cited table reports ratios for both TJ-v2-base and TJ-v2-full under these settings.
  • Full-Model Speedup: TetraJet-v2 preserves clear end-to-end speedup over TE-FP8 on full Llama2-structured models, beyond single-layer microbenchmarks.The full-model comparison is reported in Table 16.
  • Scope and Limitations: The reported acceleration focuses on model forward and backward computation rather than complete system-wide speedup.Communication, data loading, runtime scheduling, optimizer states, and other framework components remain outside the current scope.
Loading 2510.27527v3…