Technical Summary: TetraJet-v2

1. Problem Statement

Large Language Model (LLM) pre-training is increasingly prohibitively expensive, driving the need for low-precision Fully Quantized Training (FQT). While 4-bit formats like NVFP4 offer significant efficiency gains over FP8 and BF16, achieving stable, near-lossless training at this precision remains a critical challenge. Prior attempts at fully 4-bit training have suffered from severe performance degradation.

The authors identify two primary optimization bottlenecks that hinder fully FP4 training, distinct from simple quantization error:

1. Weight Oscillation: A dominant failure mode where small updates to high-precision master weights cause their quantized representations to flip repeatedly between adjacent bins (e.g., between 0 and 0.5). This phenomenon disrupts optimization, particularly as learning rates decay, preventing the model from converging to a stable local minimum.
2. Outlier Features: A small number of activation channels exhibit massive magnitudes, creating dynamic ranges that FP4 cannot represent accurately. This leads to high scaling factors that crush the precision of smaller values within the same quantization group.

Existing methods, such as NVIDIA's NVFP4 recipe, often rely on mixed-precision (keeping some blocks in BF16) or fail to fully address these optimization dynamics in a pre-training-from-scratch setting.

2. Methodology: TetraJet-v2

TetraJet-v2 is an end-to-end 4-bit FQT method designed specifically for NVFP4, applying the format to activations, weights, and gradients in all linear layers. The methodology addresses the identified bottlenecks through three core components:

A. Unbiased Double-Block Quantization for NVFP4 Linear Layers

To maximize accuracy within the numerical constraints of NVFP4 (which uses 16-element groups with 8-bit scaling), the authors propose a double-block quantization scheme:

• Structure: An outer block of 128 elements is scaled globally, while inner blocks of 16 elements are scaled locally. This avoids computing the maximum absolute value over the entire tensor, improving hardware friendliness and accuracy compared to per-tensor scaling.
• Unbiased Gradient Estimation: The linear layer design ensures unbiased gradient estimation by:
  • Aligning tensors in the backward pass (estimating $\hat{X}$ rather than $X$).
  • Using Stochastic Rounding in the backward pass to ensure $E[\text{round}(x)] = x$, preventing bias in gradient expectations.
  • Adopting a $1 \times 16$ weight group shape (non-square) for higher accuracy, requiring specific alignment for the transposed weight matrix.

B. OsciReset: Oscillation Suppression

To address weight oscillation, the authors introduce OsciReset, the first algorithm specifically designed to suppress this bottleneck in LLM pre-training.

• Detection: The method tracks the optimization trajectory of weights, calculating an "Oscillation Risk" metric defined as the ratio of the distance traveled by the quantized weight ($dist_Q$) to the distance traveled by the master weight ($dist_M$). High risk indicates a weight is flipping bins despite small master weight updates.
• Mechanism: When the learning rate is low and oscillation risk exceeds a threshold, the master weight of the oscillating element is reset to the center of its current quantization bin.
• Advantage: Unlike previous methods that freeze weights or use heavy regularization (which harm global optimization), OsciReset allows the weight to continue optimizing on a new trajectory less prone to oscillation, without degrading performance immediately.

C. OutControl: Outlier Control

To manage outliers, TetraJet-v2 employs a hybrid strategy:

• Random Hadamard Transformation (RHT): Applied to the backward pass computations ($dX$ and $dW$) to distribute outlier energy more evenly across channels, mitigating quantization variance.
• Static Outlier Retention: Recognizing that outlier channels are structurally consistent (specific channels always have high norms), the method identifies the top $p\%$ (e.g., 10%) of outlier channels during a brief calibration window. These channels are retained in higher precision (e.g., FP8) for both forward and backward passes, while the rest of the network remains in FP4. This static selection avoids the computational overhead of dynamic per-batch adjustments.

3. Key Contributions

1. Unbiased Double-Block Quantization: A novel NVFP4 linear layer design with unbiased gradient estimation and optimal block shapes ($1 \times 16$) for LLMs, outperforming existing NVFP4 recipes.
2. OsciReset: The first effective algorithm to suppress weight oscillation in LLMs during pre-training, resolving a primary source of instability in low-bit training.
3. OutControl: A mix-precision algorithm that combines RHT for the backward pass with static outlier retention for both forward and backward passes to preserve accuracy.
4. Comprehensive Evaluation: Extensive pre-training experiments on OLMo-2 models (up to 370M parameters, 212B tokens) demonstrating the efficacy of the full stack.

4. Experimental Results

The authors evaluated TetraJet-v2 on OLMo-2 models (70M, 150M, and 370M parameters) trained from scratch.

• Performance Gap Reduction: TetraJet-v2-full (incorporating all proposed methods) reduced the performance gap to full-precision (BF16) training by an average of 51.3% across model sizes.
• Perplexity (PPL): On the 370M model trained with 212B tokens, TetraJet-v2-full achieved a validation PPL of 25.11, compared to 23.70 for BF16 and 26.20 for the NVIDIA NVFP4 baseline.
• Downstream Tasks: The method achieved the highest average performance among all FP4 methods on a suite of benchmarks (ARC, BoolQ, MMLU, etc.), with a 370M model reaching 43.60% accuracy on average.
• Efficiency: CUDA kernel implementations achieved an end-to-end speedup of 1.67× over FP8 (TransformerEngine) on RTX 5090 hardware, verifying practical performance gains despite the added complexity of oscillation suppression and outlier control.
• Ablation Studies:
  • Removing OsciReset led to severe performance degradation and a rebound in loss.
  • OutControl improved accuracy in both forward and backward passes.
  • The method was shown to be robust to hyperparameter choices (e.g., oscillation threshold, start time).

5. Significance and Claims

The paper claims that TetraJet-v2 represents a significant step toward making fully low-bit (4-bit) pre-training of LLMs viable and competitive with high-precision training.

• Optimization vs. Representation: The authors emphasize that the performance gap in low-bit training is not solely due to quantization error (representation) but is critically driven by optimization issues (oscillation and outliers) that previous methods failed to address systematically.
• Generalizability: Unlike prior oscillation mitigation techniques designed for Vision Transformers (which often rely on fine-tuning or specific hyperparameter tuning), OsciReset is presented as a simple, effective method applicable to LLM pre-training from scratch with minimal tuning.
• Practicality: By achieving a 1.67× speedup over FP8 while narrowing the accuracy gap to BF16 by over 50%, the work demonstrates that fully quantized training can offer both computational efficiency and model quality, potentially reducing the hardware and energy costs of foundation model research.

The authors acknowledge limitations regarding the scale of their experiments (up to 370M parameters) and the current lack of mainstream hardware support for mixed-precision strategies like FP6 $\times$ FP4, which they explored as a potential future direction for further accuracy gains.