PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimizing Large Language Model Training Using FP4 Quantization

Ruizhe Wang, Yeyun Gong, Xiao Liu, Guoshuai Zhao, Ziyue Yang, Baining Guo, Zhengjun Zha, Peng Cheng

arXiv:2501.17116v2cs.LGcs.CL

TL;DR

Training large language models requires extensive computational resources, motivating more efficient low-precision methods while FP4 remains challenging because of quantization errors. This paper introduces an FP4 training framework with differentiable weight-gradient estimation and activation outlier compensation, achieving performance comparable to BF16 and FP8 while scaling to 13B-parameter models trained on 100B tokens.

  • Problem

    Training large language models demands substantial time, energy, and financial resources, while FP4 training is challenged by significant quantization errors and limited representational capacity.

  • Method

    The framework combines a differentiable gradient estimator for FP4 weight updates with outlier clamping and compensation for FP4 activations, using vector-wise quantization and mixed-precision training for stability.

  • Results

    FP4-trained models achieve accuracy comparable to BF16 and FP8 models, with minor training-loss gaps, scaling to 13B parameters and 100B training tokens.

  • Takeaways & Limitations

    The findings demonstrate the feasibility of FP4-based LLM training and provide a foundation for improving ultra-low-precision computing and future 4-bit hardware kernels.

  • Takeaways & Limitations

    Because dedicated FP4 Tensor Cores are unavailable, the experiments use FP4 simulations with additional casting overhead and do not extend to extremely large models or trillion-token datasets.

Abstract

from arXiv · show

The growing computational demands of training large language models (LLMs) necessitate more efficient methods. Quantized training presents a promising solution by enabling low-bit arithmetic operations to reduce these costs. While FP8 precision has demonstrated feasibility, leveraging FP4 remains a challenge due to significant quantization errors and limited representational capacity. This work introduces the first FP4 training framework for LLMs, addressing these challenges with two key innovations: a differentiable quantization estimator for precise weight updates and an outlier clamping and compensation strategy to prevent activation collapse. To ensure stability, the framework integrates a mixed-precision training scheme and vector-wise quantization. Experimental results demonstrate that our FP4 framework achieves accuracy comparable to BF16 and FP8, with minimal degradation, scaling effectively to 13B-parameter LLMs trained on up to 100B tokens. With the emergence of next-generation hardware supporting FP4, our framework sets a foundation for efficient ultra-low precision training.

1. Introduction

LLM scaling increases training resource demands, motivating lower-precision computation. This work introduces an FP4 training framework that addresses quantization errors and achieves performance comparable to higher-precision baselines.

  • Motivation: Large-scale LLM training requires substantial time, energy, financial resources, and computational infrastructure.Llama 3 405B reportedly used up to 16K H100 GPUs for 54 days, while GPT-4 required extraordinary computational power.
  • Motivation: Quantization reduces training costs by using low-bit arithmetic kernels that save memory and accelerate computation.Existing training commonly uses FP32 or FP16/BF16, while quantization enables 8-bit and 4-bit formats.
  • Research Gap: FP4 training remains an open research question despite FP8 tensor cores demonstrating training feasibility.New hardware supports low-bit kernels, including FP8 on Hopper GPUs and FP4 on Blackwell GPUs.
  • Research Gap: Direct FP4 quantization can substantially degrade accuracy because its limited dynamic range increases overflow and underflow risks.FP4 provides only 16 distinct representable values, making quantization errors especially difficult to control.
  • Approach: The proposed framework combines a differentiable weight quantization estimator with outlier clamping and compensation for activations.The activation strategy uses clamping and a sparse auxiliary matrix to preserve quantization accuracy and model performance.
  • Results: FP4-trained models achieve accuracy comparable to BF16 or FP8 models, scaling to 13B parameters and 100B training tokens with a minor training-loss gap.Zero-shot downstream results are competitive with BF16 models under the same hyperparameters.

2. Preliminaries

The preliminaries define floating-point formats by their sign, exponent, and mantissa fields, then specify the E2M1 FP4 representation and absmax quantization procedure.

  • Floating-Point Representation: Floating-point formats consist of a 1-bit sign, exponent bits, and mantissa bits, denoted ExMy.The exponents x and y indicate the numbers of exponent and mantissa bits.
  • FP4 Format: This work adopts the E2M1 format for FP4, using 2 exponent bits and 1 mantissa bit.FP16 uses E5M10, BF16 uses E8M7, and FP8 commonly uses E4M3 or E5M2.
  • Quantization: FP quantization uses uneven quantization intervals and a larger dynamic range than integer quantization.The method quantizes high-precision tensors such as FP16 to FP4 with the commonly used absmax approach.
  • Quantization: For E2M1 FP4, the maximum representable absolute value is 6.0, and quantization uses a lookup table.The lookup table is implemented in a custom CUDA kernel because FP4 supports 16 distinct values.

3. Methodology

The FP4 training scheme quantizes weights and activations for low-precision matrix multiplication while using DGE to improve weight gradients and OCC to preserve activation information. These methods address gradient vanishing, quantization error, and outlier-driven underflow through differentiated backward computation, clamping, compensation, and precision-aware processing.

  • Differentiable Gradient Estimator: DGE reduces weight-gradient estimation error by retaining hard quantization during the forward pass and adding a differentiable correction term during backpropagation.This preserves hardware-efficient forward computation while addressing the zero derivative of hard quantization and the simplifying assumption of STE.
  • FP4 Training Scheme: FP4 linear-layer training quantizes both weights and activations before GeMM, then applies two scaling factors to the result for computational correctness.The scheme converts high-precision tensors to FP4 using look-up table quantization to leverage FP4 tensor cores.
  • Differentiable Gradient Estimator: The DGE approximation sharpens as k increases, spans 14 E2M1 quantization intervals, and caps derivative magnitude at 3.0 to prevent infinite gradient spikes.The approximation is illustrated for the first positive interval [0, 0.5] and the full E2M1 dynamic range [−6.0, 6.0].
  • Outlier Clamping and Compensation: Activation outliers expand tensor dynamic range and cause most values to underflow toward zero after FP4 quantization.OCC identifies large-magnitude values using a quantile threshold and clamps them to restrict the activation range.
  • Outlier Clamping and Compensation: OCC combines clamped FP4 computation with high-precision sparse compensation, reducing quantization loss while introducing a computational-efficiency versus numerical-accuracy trade-off.The method processes the clamped matrix with FP4 GeMM and the residual with high-precision sparse matrix multiplication; lower quantile thresholds increase compensation strength and cost.

4. Experiment

The experiments evaluate FP4 training across model sizes, precision settings, downstream tasks, and ablations. FP4 closely matches BF16 overall, while DGE, OCC, and vector-wise scaling address distinct quantization challenges.

  • Implementation: The framework trains LLaMA models from scratch with 4-bit GeMM quantization, token-wise activation scaling, channel-wise weight scaling, and mixed precision for non-GeMM operations.FP4 performance is emulated with H-series GPU FP8 Tensor Cores; gradients and optimizer states use higher precisions.
  • Main Results: After 100B tokens, FP4 training losses are 2.55 versus 2.49 for BF16 at 1.3B, 2.17 versus 2.07 at 7B, and 1.97 versus 1.88 at 13B parameters.The BF16 and FP4 curves largely overlap, with FP4 showing slightly higher loss.
  • Main Results: FP4-trained models achieve comparable or slightly higher average zero-shot accuracy than BF16 models across eight downstream tasks at matched model sizes.The evaluation includes Arc, BoolQ, HellaSwag, LogiQA, PiQA, SciQ, OpenbookQA, and Lambada.
  • Main Results: FP4 models achieve comparable or slightly lower perplexity than BF16 models on Lambada, the Pile 10k, and Wikitext.Larger models also achieve lower perplexity under the same training-token budget.
  • Ablation Study: DGE improves weight-only 4-bit convergence, OCC prevents activation-only FP4 divergence, and vector-wise scaling avoids the significant errors of tensor-wise FP4 scaling.For OCC, smaller α improves accuracy but increases computational cost; the authors favor α = 0.99 for overall efficiency.

5. Related Work

Prior quantization work largely targets inference, CNNs, DNNs, or fine-tuning rather than full FP4 LLM pretraining. The paper positions its framework as the first FP4 LLM training approach validated from scratch and aligned with emerging hardware.

  • Quantized Training and Inference: Earlier 4-bit, 2-bit, and 1-bit methods primarily address LLM inference and often require calibration, compensation, parameter search, or retraining.These requirements distinguish inference quantization from the paper’s full-precision-training setting.
  • 4-bit Quantization: Existing fully quantized training studies include CNN-focused FP4 work and INT4 fine-tuning, but lack validation for full FP4 LLM pretraining.MXFP introduced quantization data for GPT-style models but did not validate full FP4 settings.
  • Differentiable Quantization: The paper’s differentiable gradient estimator differs from prior differentiable quantization methods by using a fixed quantization function and replacing STE during backpropagation.Prior approaches rely on learnable quantization parameters updated through backpropagation.
  • Handling Outliers: Existing activation-outlier methods mainly target inference, whereas this work addresses channel-specific outliers during LLM training under matrix-multiplication constraints.Channel-wise quantization can reduce quantization loss but conflicts with linear-layer computation structure.

6. Limitation

The study cannot directly measure native FP4 speed or energy gains because dedicated FP4 Tensor Cores are unavailable. Its experiments also remain bounded to models up to 13B parameters and datasets up to 100B tokens.

  • Hardware Boundary: The absence of dedicated FP4 Tensor Cores prevents direct measurement of native FP4 speedup and energy-efficiency gains.Current experiments rely on FP4 simulations that add precision-casting overhead and prolong runtime.
  • Scalability Boundary: Resource constraints leave scalability to extremely large models and trillion-token datasets untested.The authors identify broader scalability investigation as future work.

7. Conclusion

The paper introduces the first FP4 pretraining framework for modern LLMs, addressing limited dynamic range and quantization precision to reduce the gap with higher-precision training.

  • The framework combines a differentiable gradient estimator with an outlier compensation mechanism to address FP4 training challenges.These techniques target quantization precision and limited dynamic range.

Impact Statement

The work presents FP4 training as a possible route toward lower computational demand and broader access to advanced AI systems, while highlighting sustainability and responsible-adoption concerns.

  • FP4-based training may reduce computational demand, supporting energy conservation and lower carbon emissions in AI development.
  • The research calls for next-generation AI accelerators optimized for 4-bit computations, while recognizing risks from misuse and bias amplification.The passage frames these broader societal implications as challenges for responsible and equitable adoption.

A. Implementation of FP4 Quantizaiton

The implementation adopts IEEE-inspired FP4 representations, selects E2M1 as a balance between range and precision, and uses lookup-table quantization in a custom CUDA kernel.

  • Floating-point representations use sign, exponent, and mantissa components, with limited mantissa bits introducing rounding errors.The representation is described using the ExMy notation for exponent and mantissa widths.
  • FP4 uses the E2M1 format because it balances dynamic range and quantization precision.Increasing exponent bits expands dynamic range, while increasing mantissa bits improves quantization-interval precision.
  • Table 4 and Figure 7 enumerate representable values across FP4 formats.The comparison supports selecting among formats according to dynamic range and quantization precision.
  • FP4 represents 16 distinct values, motivating lookup-table quantization in a custom CUDA kernel.The kernel applies element-wise quantization operations that can be parallelized on GPUs.

B. Theoretical Analysis on Speed Performance and Overhead

The theoretical analysis estimates FP4 acceleration from Transformer-layer FLOPs and then accounts for DGE and OCC overheads, while deriving DGE gradients under vector-wise scaling.

  • Theoretical speedup: For a representative 7B model with h = 4096 and s = 2048, the ideal theoretical FP4 speedup is 3.12.The estimate compares forward and backward computation in FP32 and FP4 precision.
  • Overhead analysis: DGE adds approximately 8 FLOPs per input element through an additional nonlinear function in the backward pass.The overhead occurs once per forward-backward iteration.
  • Overhead analysis: OCC adds sparse FP8 matrix multiplications for activation outlier compensation, using α = 0.99 to maintain high sparsity.The sparse compensation operations are applied to the four primary GEMMs in a Transformer block.
  • Adjusted speedup: DGE and OCC overheads are estimated at 0.1% and 5.6%, reducing the ideal speedup from 3.12 to 2.95.The paper characterizes this reduction as an acceptable trade-off between computational efficiency and model accuracy.
  • DGE derivation: DGE applies a correction term to weight gradients while preserving hard quantization in the forward computation.The derivation accounts for channel-wise scaling factors and element-wise operations.
  • DGE derivation: Under vector-wise quantization, scaling and un-scaling cancel element-wise, yielding the simplified true weight-gradient expression.The final expression matches the main-body formulation when the DGE correction uses scaled weights.

C.3. Mathematical Soundness of DGE Clipping

DGE clipping addresses an unbounded derivative near x = δ/2 that can destabilize training. Smoothing with epsilon bounds the correction term, while clipping provides an equivalent practical stabilization.

  • In practice, the derivative f′(x) is clipped at 3.0 to prevent instability during training.
  • The DGE derivative can diverge near x = δ/2 because a denominator vanishes, potentially causing gradient explosions during optimization.
  • Replacing the sharp absolute value with a smooth approximation removes the singularity while preserving differentiability and numerical stability.
  • The epsilon-regularized derivative bounds the DGE correction term.
  • Clipping the DGE correction is mathematically equivalent to smoothing because k and ϵ are constants, while both approaches bound gradient magnitude near x = δ/2.

D. Analyzing Quantization Difficulty Through Tensor Distribution

Weight and activation tensors differ substantially in their quantization difficulty: activations have wider and more irregular distributions, with outliers concentrated in specific channels. These properties make activation quantization the central challenge examined by the distribution analysis.

  • Figures 8–10 show weight distributions, while Figures 11–13 show activation distributions from 30,000 iterations of LLaMA 1.3B training.The figures use a logarithmic y-axis for visualization.
  • Weight tensors generally have smaller dynamic ranges, whereas activation tensors have significantly larger ranges that make quantization more difficult.
  • Weights typically follow near-normal distributions with small outliers, while layer-normalization and transformer-layer activations often have irregular distributions with numerous outliers.
  • Activation outliers tend to occur in specific channels, as illustrated by a heatmap of GeLU outputs from the first transformer layer.
  • The analysis identifies effective handling of activation outliers as critical for FP4 quantization and motivates further study of their distributions and behavior.
Loading 2501.17116v2…