DEV Community

Prabhakar Chaudhary
Prabhakar Chaudhary

Posted on

Stable FP4 Pretraining with Transpose-Invariant 2D Block Scaling

Stable FP4 Pretraining with Transpose-Invariant 2D Block Scaling

Billing Support — September 03, 2026

A tensor value can receive one scale during the forward pass and a different scale after the same tensor is transposed for backpropagation. At FP4 precision, that mismatch is not a minor rounding detail: it changes the quantized representation used to calculate gradients and can introduce systematic bias into parameter updates.

The September 3 cs.LG work on stable FP4 pretraining through block scaling addresses this failure mode by replacing transpose-sensitive one-dimensional scaling groups with square two-dimensional blocks. The broader recipe combines that structural change with truncation-free scaling, stochastic rounding, deterministic Hadamard rotations, and selective BF16 computation.

Why 1D microscaling breaks under transposition

FP4 offers limited dynamic range and relatively large quantization error. Microscaling formats such as MXFP4 or NVFP4 compensate by dividing a tensor into small groups and assigning each group its own scale.

For a one-dimensional group containing values x, quantization can be summarized as:

q(x, s) = round(x / s)

where s is derived from the values in that group. The problem is that matrix transposition changes which values belong together.

Suppose a matrix is grouped along rows during a forward matrix multiplication. Backpropagation may require its transpose, causing the corresponding computation to group values along what were previously columns. The numerical values have not changed, but their neighbors—and therefore their group scales—have.

The result is scale inconsistency:

q(X, row-scales)^T != q(X^T, row-scales-of-X^T)

This means the forward and backward passes operate on different low-precision approximations of the same underlying tensor. According to the paper’s reported analysis, this inconsistency produces biased gradients. Deterministic rounding cannot repair the structural problem because it only decides how values map to levels after the incompatible scales have already been selected.

Square blocks preserve scale assignments

The proposed fix partitions matrices into square two-dimensional blocks, with 32×32 given as an example. Each block receives a shared scale.

When the matrix is transposed, every square block is also transposed, but its membership remains intact: the same values stay together. A block at one matrix coordinate moves to its transposed coordinate without being regrouped into unrelated sets.

Conceptually:

scale(B) = scale(B^T)

provided the scale calculation is itself insensitive to element order. Forward and backward matrix multiplications can therefore reuse consistent quantized representations instead of deriving scales from different one-dimensional slices.

This does not eliminate FP4 error. It removes one specific source of systematic error: transpose-induced reassignment of values to scaling groups. That distinction matters for implementers. Two-dimensional scaling is the structural foundation, while the remaining techniques control clipping, rounding error, and outliers.

The rest of the stability recipe

Truncation-free scaling

A scale that is too small pushes large values outside FP4’s representable range. Those values are clipped, creating systematic distortion.

Truncation-free scaling chooses the scale so all values in the block fit within the available range. This avoids clipping rather than accepting it as an ordinary quantization effect. The trade-off is that a single large value may force a coarser quantization step for every other value in the block.

That trade-off explains why scaling alone is insufficient: preserving the largest value can reduce resolution for the majority of smaller values.

Stochastic rounding

Round-to-nearest deterministically maps a value to one adjacent representable level. Under repeated low-precision operations, those choices can accumulate directionally.

Stochastic rounding instead selects between neighboring levels probabilistically so that the expected quantized value equals the original value. In compact form:

E[q(x)] = x

This property is especially relevant to gradients, where persistent directional error can alter optimization. Stochastic rounding does not make an individual quantization exact; it targets unbiased behavior in expectation.

Deterministic Hadamard rotations

Outliers create another tension. If a few coordinates carry unusually large magnitude, truncation-free scaling must accommodate them, leaving fewer useful levels for the rest of the block.

Hadamard rotations in low-precision training redistribute outlier energy across coordinates before quantization. The evidence reports that deterministic Hadamard transforms were more effective than randomized variants for stabilizing the complete pipeline, particularly when quantizing weight gradients.

The rotation does not discard information. Its role is to produce a representation whose magnitudes are easier to quantize with one block scale.

Selectively retained BF16 paths

The method is not a claim that every operation should run in FP4. Some paths remain numerically sensitive, especially attention operations involving softmax and dot-product interactions.

The reported mixed-precision policy quantizes dense Q, K, and V projections while retaining sensitive attention paths in BF16. Related full-stack recipes also protect selected compact subspaces in BF16 while executing dominant dense operations in FP4.

For engineers, the practical principle is selective precision: use FP4 where tensor operations dominate cost, but retain BF16 where quantization error would be amplified by the operation.

Interpreting parity and speed claims

The reported evidence indicates performance near BF16 baselines, often with less than 1.5% perplexity degradation when the stabilization techniques are applied together. “Parity” should therefore be read as comparable training quality under the evaluated configurations, not numerical equivalence at every step.

The broader FP4 training literature also reports speedups as high as 4.64× in specific rollout-heavy tasks. That figure is conditional, not a universal pretraining multiplier. Real gains require hardware with native FP4 tensor support and workloads where reduced arithmetic and memory-bandwidth pressure affect end-to-end runtime. BF16 fallbacks, rotations, scale calculation, and data conversion all consume part of the theoretical saving.

Engineering implications and limits

A practical implementation needs more than changing a datatype:

  • Quantization metadata must represent two-dimensional block scales.
  • Matrix layouts and kernels must preserve block identity across transposes.
  • Forward, activation-gradient, and weight-gradient paths need coordinated policies.
  • Stochastic rounding must be integrated into training kernels.
  • Hadamard transforms add operations that must be measured against their stability benefit.
  • Attention and other sensitive paths require explicit BF16 exceptions.

The main limitation is complexity. Square block scaling changes kernel design and metadata handling, while mixed precision creates more execution paths to validate. Truncation-free scaling prevents clipping but cannot prevent an outlier from reducing effective resolution. Stochastic rounding is unbiased only in expectation, and the reported results do not imply that every model, optimizer, or hardware stack will match BF16.

The core result is narrower and more useful: if FP4 groups are not preserved under the transposes required by backpropagation, scale inconsistency can bias gradients. Square 2D blocks preserve those assignments. Combined with careful rounding, outlier redistribution, clipping avoidance, and selective BF16 retention, they provide a technically credible route to stable FP4 pretraining.

Top comments (0)