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

TL;DR

TetraJet-v2 introduces a novel 4-bit fully-quantized training method for large language models, effectively addressing oscillation and outliers to enable efficient low-precision training with minimal performance loss.

Contribution

The paper presents TetraJet-v2, a comprehensive 4-bit training approach that includes new algorithms for quantization, oscillation suppression, and outlier control, advancing low-precision LLM training.

Findings

Outperforms prior FP4 methods on models up to 370M parameters.

Reduces performance gap to BF16 by 51.3% on average.

Achieves 1.67x speedup over FP8 training.

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.