HyperAIHyperAI

Command Palette

Search for a command to run...

LLM FP4 사전학습에서 수축 편향 재고찰: 기하학적 기원, 체계적 영향, 및 UFP4 레시피

초록

FP4 정밀도 훈련은 대규모 언어 모델(LLM) 사전 훈련 과정에서 메모리 및 연산 비용의 현저한 감소를 약속하지만, 현재 NVIDIA Blackwell/Rubin급 시스템과 AMD MI350 시리즈 GPU를 포함한 FP4 하드웨어 경로 및 훈련 레시피는 E2M1 데이터 요소에 중심을 두고 있습니다. 본 연구에서는 이러한 선택의 근본적인 한계를 규명합니다. E2M1과 같은 비균일 형식은 표현 가능한 구간의 기하학적 비대칭으로 인해 발생하는 체계적인 음의 반올림 오차인 수축 편향(Shrinkage Bias)을 본질적으로 겪습니다. 본 연구는 이러한 편향이 레이어 간에乗法的으로 누적되며, 무작위 하달드 변환(RHT)에 의해 증폭됨을 보여줍니다. 이는 기존 E2M1 기반 FP4 레시피에서 관찰된 훈련 불안정성에 대한 통합적인 설명을 제공합니다.반면, 균일 그리드(E1M2/INT4)는 이러한 그리드 기하학적 오류를 피하고, RHT로 인해 개선된 버킷 활용도를 더 높은 양자화 품질로 전환합니다. 이러한 발견을 바탕으로, 본 연구는 세 가지 훈련 GEMM 모두에 RHT를 적용하고 확률적 반올림(stochastic rounding)은 dY에만 국한시키는 균일 4비트 훈련 레시피인 UFP4를 제안합니다. Dense 1.5B, MoE 7.9B, 그리고 MoE 124B에서의 장기 사전 훈련 실험에서 UFP4는 스케일링 법칙 분석 및ablative 연구(a ablation studies)의 지지를 받으며, 강력한 E2M1 기반 베이스라인 대비 BF16 상대 손실 분해가 consistently 낮음을 보여줍니다. 본 연구의 결과는 향후 가속기가 E2M1과 함께 E1M2/INT4 스타일의 균일 4비트 그리드를 일차원 훈련 원리(first-class training primitives)로 지원해야 함을 시사합니다.

One-sentence Summary

Ant Group researchers identify Shrinkage Bias in E2M1-based FP4 training as a systematic negative rounding error arising from geometric bin asymmetry and propose UFP4, a uniform 4-bit recipe using E1M2/INT4 grids with the Random Hadamard Transform applied to all three training GEMMs and stochastic rounding restricted to dY, achieving lower BF16-relative loss degradation across Dense 1.5B, MoE 7.9B, and MoE 124B pretraining.

Key Contributions

  • The paper identifies Shrinkage Bias, a systematic negative rounding error caused by the geometric asymmetry of the non-uniform E2M1 grid, and shows that it accumulates multiplicatively across layers and is amplified by the Random Hadamard Transform, explaining the training instability observed in existing E2M1-based FP4 recipes.
  • It proposes UFP4, a uniform 4-bit training recipe that applies RHT to all three training GEMMs, restricts stochastic rounding to the dY gradient path, and uses an E1M2/INT4-style uniform grid to convert improved bucket utilization from RHT into higher quantization quality.
  • On Dense 1.5B, MoE 7.9B, and MoE 124B long-run pretraining, UFP4 consistently achieves lower BF16-relative loss degradation than strong E2M1-based baselines, supported by scaling-law analysis and ablation studies.

Introduction

The authors investigate the training of Large Language Models using 4-bit floating-point formats, a technique that promises significant reductions in memory and computational cost. Current hardware and training recipes predominantly rely on the non-uniform E2M1 format, but end-to-end FP4 pretraining remains unstable and suffers from accuracy loss compared to higher precision baselines. Prior work attempts to mitigate this instability with techniques like the Random Hadamard Transform (RHT) and stochastic rounding, yet retains the problematic E2M1 data format.

The authors identify a fundamental limitation: non-uniform grids like E2M1 inherently suffer from Shrinkage Bias, a systematic negative rounding error caused by the geometric asymmetry of their representable bins. They demonstrate that this bias accumulates multiplicatively across network layers, causing signal decay, and is paradoxically amplified by RHT, which shifts tensor mass into the format's most asymmetric regions. The authors’ main contribution is UFP4, a training recipe that replaces the E2M1 grid with a uniform 4-bit format like E1M2 or INT4. This approach bypasses the grid-geometry error, enabling the effective use of RHT on all critical matrix multiplications while restricting stochastic rounding to a single gradient tensor, thereby achieving superior training stability and accuracy at scale.

Dataset

The authors describe the quantization formats and the blockwise quantization procedure used to construct low-precision training data.

  • Format codebooks: Two 4-bit floating-point formats, E2M1 and E1M2, and one 4-bit integer format (INT4) are considered. Each format defines a normalized codebook GGG with a fixed set of representable levels.
  • Blockwise quantization: A tensor is partitioned into contiguous blocks. For each block, a shared scale sBs_BsB is computed from the maximum absolute value in the block divided by the largest magnitude level in the codebook. Every element is then scaled and mapped to a codebook level via a rounding rule.
  • Rounding rules: The paper supports both Round-To-Nearest-Even (RTNE), which is deterministic, and Stochastic Rounding (SR), which samples a level with probability proportional to the distance between adjacent levels. RTNE is noted as the widely adopted choice in recent low-precision recipes such as the NVFP4 method.
  • Usage in the model: The quantized tensor QG(T)Q_G(\mathbf{T})QG(T) is used throughout the training pipeline; in arithmetic expressions it denotes the dequantized numerical tensor, ensuring that the low-precision representation is directly integrated into forward and backward passes.

Method

The authors address the challenge of quantizing training tensors, which often contain outlier coordinates that cause most codebook levels to be underutilized. To mitigate this, they leverage Random Hadamard Transforms (RHT). RHT applies a norm-preserving rotation that disperses outlier energy across all coordinates before quantization. The authors use Sylvester Hadamard matrices defined recursively as:

H2n=12[HnHnHnHn],H1=[1],\mathbf { H } _ { 2 n } = \frac { 1 } { \sqrt { 2 } } \left[ \begin{array} { l l } { \mathbf { H } _ { n } } & { \mathbf { H } _ { n } } \\ { \mathbf { H } _ { n } } & { - \mathbf { H } _ { n } } \end{array} \right] , \qquad \mathbf { H } _ { 1 } = [ 1 ] ,H2n=21[HnHnHnHn],H1=[1],

where HnHn=In\mathbf{H}_{n}^{\top}\mathbf{H}_{n}=I_{n}HnHn=In for n=2kn=2^{k}n=2k. An RHT additionally applies a random sign matrix Sn=diag(ϵ1,,ϵn)\mathbf{S}_{n}=\operatorname{diag}(\epsilon_{1},\ldots,\epsilon_{n})Sn=diag(ϵ1,,ϵn), ϵi{1,+1}\epsilon_{i}\in\{-1,+1\}ϵi{1,+1}, before the Hadamard transform. Since Hn=SnHn\mathbf{H}_{n}^{\prime}=\mathbf{S}_{n}\mathbf{H}_{n}Hn=SnHn is orthogonal, applying the same transform to the shared GEMM dimension preserves the full-precision result:

Hn=SnHn,Y=XW=(XHn)(WHn).\begin{array} { r } { \mathbf { H } _ { n } ^ { \prime } = \mathbf { S } _ { n } \mathbf { H } _ { n } , } \\ { \mathbf { Y } = \mathbf { X } \mathbf { W } ^ { \top } = ( \mathbf { X } \mathbf { H } _ { n } ^ { \prime } ) ( \mathbf { W } \mathbf { H } _ { n } ^ { \prime } ) ^ { \top } . } \end{array}Hn=SnHn,Y=XW=(XHn)(WHn).

The corresponding low-precision GEMM quantizes the rotated operands:

Y^=QG(χHn)QG(WHn).\widehat { \mathbf { Y } } = Q _ { G } ( \boldsymbol { \chi } \mathbf { H } _ { n } ^ { \prime } ) \, Q _ { G } ( \mathbf { W } \mathbf { H } _ { n } ^ { \prime } ) ^ { \top } .Y=QG(χHn)QG(WHn).

While RHT reduces outlier dominance, its effectiveness depends heavily on the underlying quantization grid. The authors identify a fundamental issue with the mainstream E2M1 format, which they term Shrinkage Bias. This bias arises from the geometric asymmetry of Round-to-Nearest-Even (RTNE) rounding bins in non-uniform grids. For an interior quantization level qiq_iqi, the RTNE rounding bin has left and right widths i\ell_{i}i and rir_{i}ri. If the density inside the bin is locally uniform, the conditional expected error is:

E[ρG(t)ttBi]=iri2=2qiqi1qi+14.\mathbb { E } [ \rho _ { G } ( t ) - t \mid t \in \mathcal { B } _ { i } ] = \frac { \ell _ { i } - r _ { i } } { 2 } = \frac { 2 q _ { i } - q _ { i - 1 } - q _ { i + 1 } } { 4 } .E[ρG(t)ttBi]=2iri=42qiqi1qi+1.

An asymmetric bin with ri>ir_{i} > \ell_{i}ri>i inherently yields a negative expected error, causing magnitude shrinkage. As shown in the figure below, the E2M1 grid exhibits these asymmetric bins at spacing-transition points, whereas uniform grids like E1M2 and INT4 maintain symmetric bins (i=ri\ell_{i}=r_{i}i=ri), thereby eliminating this geometric source of bias.

To resolve this, the authors propose UFP4, a 4-bit training recipe based on uniform grids (E1M2/INT4). Once RHT shifts tensors from a dynamic-range-limited to a local-resolution-limited regime, the 4-bit grid prioritizes local magnitude preservation over extreme dynamic range. Refer to the framework diagram for the complete training pipeline.

In the UFP4 recipe, for every linear-layer GEMM, the authors apply RHT and quantize the operands to the uniformly spaced grid. Crucially, while E2M1 recipes typically restrict RHT to the weight-gradient path to avoid compounding geometric errors, the unbiased nature of UFP4 allows them to safely enable RHT across all three GEMMs: FPROP (fwd_y), DGRAD (bwd_dx), and WGRAD (bwd(dw)). Stochastic rounding (SR) is applied to dYdYdY to preserve gradient expectations. The scale hierarchy design remains orthogonal to this format-level solution, allowing for single-level, two-level, or block scales depending on hardware efficiency.

Experiment

The experiments examine how FP4 grid geometry and random Hadamard transform (RHT) scope interact across local quantization and end-to-end training. Single-tensor and GEMM-output diagnostics show that RHT reverses the preferred format ranking on outlier-heavy tensors, making the uniform E1M2 grid outperform E2M1 after rotation. In long-run pretraining across dense and MoE models up to 124B parameters, the E1M2-based UFP4 recipe consistently reduces the BF16-relative loss gap compared to an E2M1 reference, and scaling-law validation confirms this advantage persists across model scales. Ablation studies reveal that full RHT coverage and stochastic rounding on the output gradient both contribute to the gain, while range-restricted E2M1 cannot substitute for native uniform-grid support, and kernel benchmarks show that fusing RHT into FP4 quantization incurs only modest overhead.

The authors conduct an ablation study to evaluate the impact of Random Hadamard Transform scope and stochastic rounding on FP4 E1M2 training. Results indicate that applying full transform coverage across all GEMM outputs yields the lowest language modeling loss, and adding stochastic rounding further improves performance. Full transform coverage outperforms partial applications, demonstrating that rotating all GEMM outputs is beneficial under the E1M2 grid. Stochastic rounding provides an additional loss reduction when combined with full transform coverage. The UFP4 recipe achieves the best overall performance compared to configurations with limited scope or no stochastic rounding.

The ablation study examines how the scope of the Random Hadamard Transform and the use of stochastic rounding affect FP4 E1M2 training. Applying the transform to all GEMM outputs consistently yields lower language modeling loss than partial coverage, confirming that full rotation is beneficial under this data format. Adding stochastic rounding on top of full transform coverage further reduces loss, and the combined UFP4 recipe achieves the best overall performance.


AI로 AI 구축

아이디어에서 출시까지 — 무료 AI 코코딩, 즉시 사용 가능한 환경, 최적의 GPU 가격으로 AI 개발을 가속화하세요.

AI 협업 코딩
바로 사용 가능한 GPU
최적의 가격

HyperAI Newsletters

최신 정보 구독하기
한국 시간 매주 월요일 오전 9시 에 이번 주의 최신 업데이트를 메일로 발송합니다
이메일 서비스 제공: MailChimp