PaperlayerView exact revision on arXiv

Source-linked AI summary

UE5M3 FP4 Block Scaling for Stable Language Model Pretraining

Robert Hu, Carlo Luschi, Paul Balanca

arXiv:2609.02846v1cs.LG

TL;DR

FP4 pretraining is challenging because E2M1 has a narrow magnitude range, while existing stabilization methods add work outside FP4 matrix multiplications. The paper uses wider UE5M3 block scales with periodic tensor scaling and a simpler recipe, then evaluates long-run Nemotron-H 8B pretraining. The proposed path finishes below native Transformer Engine NVFP4 in held-out NLL and has higher downstream point estimates on all three reported aggregates, while a native ablation measures 21.2% higher model-body throughput.

  • Problem

    FP4 pretraining must preserve optimization while repeatedly quantizing changing activations, weights, and gradients despite E2M1's narrow magnitude range and added stabilization work.

  • Method

    The paper combines E2M1 payloads with UE5M3 block scales, periodic sample-and-hold tensor scaling, 2D weight scaling, selective upstream-gradient stochastic rounding, no RHT, and FP4 in all eligible internal linears.

  • Results

    The proposed block-16 path finishes below native Transformer Engine NVFP4 in held-out NLL and exceeds it on Core 9, OLMES MMLU, and MMLU-Pro MC point estimates.

  • Takeaways & Limitations

    The results demonstrate end-to-end software-emulated UE5M3 FP4 pretraining with a simpler recipe and motivate native support for UE5M3 block scaling.

Abstract

from arXiv · show

Stable 4-bit floating-point (FP4) pretraining is difficult because the E2M1 payload represents only a narrow range of magnitudes. NVIDIA's Transformer Engine \nv{} recipe addresses this with current-tensor scaling, a randomized Hadamard transform (RHT), and bfloat16 (BF16) final layers, adding work outside the FP4 matrix multiplications. We instead pair E2M1 payloads with unsigned E5M3 (\ue{}) block scales. Their wider range permits periodic tensor scaling, while our recipe applies selective stochastic rounding to backward gradients, omits RHT, and uses FP4 in all eligible internal linears. We pretrain a Nemotron-H 8B model for nearly 190 billion tokens. Compared with Transformer Engine \nv{}, the proposed block-16 recipe finishes with lower final-window training loss and, under their respective quantized-inference policies, lower validation loss measured as held-out negative log-likelihood. Its quantized-inference downstream point estimates are also higher on all three reported aggregates. A native \nv{} execution ablation that jointly removes RHT and the BF16 final-block exemption increases measured model-body token throughput by 21.2\%. These results demonstrate end-to-end software-emulated \uefp{} pretraining with a simpler recipe and motivate native support for \ue{} block scaling.

A Simpler FP4 Pretraining Recipe with Delayed Tensor Scaling

The section foregrounds wider block-scale range, a proposed FP4 pretraining recipe, and a software model of native execution.

  • Wider block-scale range is identified as a central theme.
  • The work presents an FP4 pretraining recipe.
  • The section also concerns a software model of native execution.

1 Introduction

FP4 pretraining must preserve optimization under a narrow E2M1 payload, motivating wider block scales and a simpler recipe. The paper introduces UE5M3 scaling and evaluates it through long-run training, deterministic controls, range analysis, and throughput measurement.

  • FP4 pretraining is difficult because E2M1 has a narrow magnitude range and requires scaling for activations, weights, and gradients.
  • The paper asks whether wider block-scale dynamic range can stabilize FP4 training with a simpler recipe.
  • The proposed recipe combines UE5M3 block scales, periodic sample-and-hold tensor scaling, 2D weight scaling, selective upstream-gradient stochastic rounding, omitted RHT, and FP4 across 112 eligible internal linears.
  • 188.7 billion tokens are used to train one seed-42 Nemotron-H 8B trajectory per configuration, with held-out NLL and downstream evaluation across checkpoints.
  • The paper derives roughly 137× default growth headroom and 30× headroom for the targeted T = 2048 override.
  • A software FP4 GEMM model matches tested native behavior in deterministic controls, while the joint no-RHT/all-linears ablation records 21.2% higher model-body token throughput.

2 Related Work

Prior work explores diverse low-precision formats, rounding strategies, transforms, and module-specific precision choices for FP4 training. This paper is positioned within efforts to make major training GEMMs quantized while addressing range, bias, and stability constraints.

  • Earlier 4-bit studies identify gradient range and quantization bias as central constraints.
  • LLM studies evaluate mixed precisions, simulated E2M1 quantization, differentiable estimators, sparse outlier compensation, and stage-specific precision.
  • Recent methods target fully quantized forward and backward GEMMs using selective stochastic rounding, RHT, or alternative unbiased microscaling estimators.
  • The proposed recipe differs by using UE5M3 scales, nearest-even saved activations, selective upstream-gradient stochastic rounding, periodic tensor references, and no BF16 backward/update phase.
  • Related work also covers native AMD MXFP4 experiments, HiFloat4 training, optimizer and attention quantization, and UE5M3 range–precision analysis.

3 Background: FP4 Microscaling

FP4 microscaling combines narrow E2M1 payloads with block-level scales to track local magnitudes, while scale formats trade range against precision. NVFP4 adds tensor scaling and several stabilization mechanisms to manage outliers and limited scale range.

  • Floating-point formats trade exponent-driven range against fraction-driven representational density.
  • Quantization rounds high-precision values into a format, with saturation or overflow for oversized values and underflow to zero for undersized nonzero values.
  • Microscaling assigns each block a scale so the narrow E2M1 payload follows local rather than tensor-wide magnitude.
  • MXFP4 uses 32 E2M1 values with E8M0 scales, whereas NVFP4 uses 16 values with finer E4M3 scales plus an FP32 tensor scale.
  • NVFP4 manages outliers and scale limitations with RHT, 2D weight scaling, stochastic rounding, BF16 final layers, and tensor scaling.
  • The evaluated baseline uses current-tensor scaling with D = 1, while D denotes optimizer steps between tensor-maximum refreshes.

4 Unsigned E5M3 Block Scales

UE5M3 repurposes the unused sign bit of a nonnegative block scale as a fifth exponent bit, expanding usable scale range without adding fraction precision. This wider codebook better covers varied block magnitudes than signed E4M3.

  • 4.1 Reallocating the unused sign bit: UE5M3 uses five exponent bits and three fraction bits, whereas E4M3 uses one sign, four exponent, and three fraction bits.Because block scales are nonnegative, UE5M3 reallocates the sign bit rather than adding a fraction bit.
  • 4.1 Reallocating the unused sign bit: UE5M3 provides roughly twice as many useful nonnegative scale encodings by eliminating codes for negative values.Both formats retain three fraction bits and therefore the same spacing at a fixed exponent.
  • 4.1 Reallocating the unused sign bit: The maximum UE5M3 scale is about 137× larger, while its smallest nonzero value is 256× smaller than finite E4M3.The smallest normal UE5M3 value is 2^-14 and the smallest subnormal is 2^-17.
  • 4.1 Reallocating the unused sign bit: UE5M3 can represent both small and large block scales directly, reducing pressure to move a shared tensor-wide scale when blocks contain outliers.The wider interval supports block-scale coverage across more heterogeneous tensor magnitudes.

5 Training Method

The method combines periodic sample-and-hold tensor references with a tunable UE5M3 target, selective stochastic rounding, and targeted higher scaling for late-layer gradients. A matched smaller-model control supports the targeted override, while the checkpoint analysis is explicitly limited to one held-out snapshot.

  • 5.1 Periodic refresh of tensor maxima: Periodic scaling samples and caches each operand’s tensor maximum, reusing that reference for D − 1 subsequent steps.The implementation uses D = 50, separate caches per activation, weight, and gradient operand, and updates the tensor-wide maximum only on refresh steps.
  • 5.2 Treating NVIDIA’s fixed 448 scale as a tunable target: The target T shifts the fixed UE5M3 codebook: larger T protects smaller block scales, while smaller T preserves more stale-growth headroom.T changes codebook placement rather than the E2M1 payload or within-band spacing.
  • 5.2 Treating NVIDIA’s fixed 448 scale as a tunable target: At T = 448, UE5M3 permits roughly 137× stale growth before saturation; T = 2048 moves small block scales about 4.6× farther from zero.The larger target trades some stale-growth headroom for underflow margin.
  • 5.2 Treating NVIDIA’s fixed 448 scale as a tunable target: The 8B recipe uses T = 2048 only for dY in the weight-gradient GEMMs of layers 45, 47, 49, and 51.Those modules retain 30× stale-growth headroom because 61,440/2,048 = 30.
  • 5.3 Checkpoint snapshot: The checkpoint snapshot found no pre-repair zero-rounded or saturated scale codes for either target, so it illustrates codebook placement rather than observed underflow or saturation.It covers one held-out sequence and does not represent the training-wide D = 50 cache history.
  • 5.2 Treating NVIDIA’s fixed 448 scale as a tunable target: The matched 350M control reduced the final-250-update mean training loss from 2.90904 to 2.90248, a difference of −0.00656, under the targeted T = 2048 override.Both runs used the same source revision, seed 42, 10,000 updates, and D = 50.
  • 5.2 Treating NVIDIA’s fixed 448 scale as a tunable target: Selective stochastic rounding uses an unbiased randomized choice between adjacent representable values, avoiding deterministic removal of values below a rounding threshold.The paper cautions that implementations using too few random bits can introduce bias.

6 Modeling Observed Native NVFP4 GEMM Outputs

The paper models observed native NVFP4 GEMM outputs using grouped accumulation, round-toward-zero cross-group additions, and optional product-lattice canonicalization. The resulting probe-matched emulator reproduces native behavior closely enough to achieve exact end-to-end gradient parity in a 1.2B control.

  • Observed native behavior: Floating-point reduction order and rounding can change GEMM outputs even when decoded FP4 operands and scales are identical.The difference arises because intermediate sums are rounded during accumulation.
  • Product-lattice canonicalization: The emulator applies ties-to-even canonicalization to multiples of 1/1024 before applying the encoded tensor-scale product α.The implementation uses torch.round(1024 * c) / 1024, with the scale product applied afterward.
  • Observed native behavior: The probe-matched FP4 emulator uses one BF16 dot product per 64-wide slice, FP32 partial sums, round-toward-zero additions, and optional product-lattice canonicalization.These rules empirically model the tested Blackwell/Transformer Engine stack.
  • Product-lattice canonicalization: 1/1024 is the coarsest tested final-grid denominator preserving all 258 native permutation matches; coarser grids destroy native bins.The 1/512 grid loses all 37 occurrences of odd bin −3371, while finer grids and no snap match this corpus.
  • End-to-end parity test: The complete 1.2B control achieves exactly zero loss difference and zero gradient relative L2 and maximum absolute error between native and emulated paths.All 1,291,929,600 compared gradient values are covered, including quantized and high-precision modules.

7 Experiments

The experiments compare a proposed UE5M3 FP4 recipe with Transformer Engine NVFP4 on matched Nemotron-H 8B pretraining and quantized inference evaluations. The proposed block-16 configuration improves reported training loss, validation NLL, downstream point estimates, and measured throughput in the relevant ablation.

  • 7.1 Setup: The study trains Nemotron-H 8B configurations for 188.7 billion tokens while matching NVIDIA’s disclosed architecture, optimizer, schedule, sequence, batch, and precision placement where possible.The study substitutes a fixed OLMo-family data mixture and uses a shorter horizon because NVIDIA’s underlying data details are unavailable.
  • 7.1 Setup: Each configuration has one seed-42 trajectory, so reported loss and downstream-score differences are descriptive rather than estimates of seed variability or statistical significance.This limits inferential claims from the comparisons.
  • 7.3 Pretraining quality and stability: All four UE5M3 FP4 trajectories have lower final-window means and endpoint losses than the Transformer Engine NVFP4 trajectory during the final 5,000 steps.The proposed block-16 recipe reaches 2.3090 versus 2.3474 for native Transformer Engine NVFP4; block-32 reaches 2.3241.
  • 7.3 Pretraining quality and stability: 205 losses above 3 and 89 gradient norms above 1 occur in the native NVFP4 no-RHT/all-linears ablation, versus none above either threshold after step 2,500 for the proposed recipe.The comparison concerns complete recipes that also differ in scale lifecycle and GEMM numerics.
  • 7.4 Validation loss: At step 30,000, proposed block-16 UE5M3 FP4 reaches 2.32230 versus 2.32592 NLL for native Transformer Engine NVFP4, a difference of −0.00362 NLL.The proposed path is lower at 8 of 12 checkpoints, including every checkpoint from step 22,500 onward.
  • 7.5 Downstream evaluation: Relative to native Transformer Engine NVFP4, proposed block-16 UE5M3 FP4 is higher by 0.13, 1.01, and 0.01 percentage points on Core 9, OLMES MMLU, and MMLU-Pro MC.These are the three reported quantized-inference aggregates.
  • 7.6 Native NVFP4 execution ablation: 21.2% higher measured model-body token throughput results from jointly removing RHT and using FP4 in the otherwise BF16-exempt final-block projections.Throughput rises from 3,212 to 3,894 tokens per second in the native NVFP4 execution ablation.

8 Conclusion

The work demonstrates end-to-end software-emulated UE5M3 FP4 pretraining with a simpler recipe than Transformer Engine NVFP4. The proposed path achieves lower training and validation loss, higher downstream point estimates, and a 21.2% native-execution throughput increase in the joint ablation.

  • 188.7 billion tokens were used to pretrain a Nemotron-H 8B model with the proposed UE5M3 FP4 recipe.
  • The recipe combines periodic tensor scaling, two-dimensional weight scaling, upstream-gradient stochastic rounding, no RHT, and FP4 in all eligible internal linears.
  • The proposed recipe achieves lower final-window mean training loss than NVIDIA’s Transformer Engine NVFP4 recipe.
  • 2.32230 versus 2.32592 is the final held-out NLL for the proposed block-16 path and native Transformer Engine NVFP4, respectively.
  • 0.13, 1.01, and 0.01 percentage points separate the proposed block-16 path from native Transformer Engine NVFP4 on Core 9, OLMES MMLU, and MMLU-Pro MC.
  • 21.2% higher measured model-body token throughput results when native execution jointly disables RHT and uses FP4 in all eligible linears.

Generative-AI Disclosure

This section specifies the quantized-inference evaluation setup and reports that activation-scale policy choice has little effect on the proposed checkpoint’s ordering at step 30,000.

  • Each loaded weight tensor’s amax is measured once, while its tensor-wide reference remains fixed across inference policies.
  • The delayed cache starts from the first inference batch rather than restoring training state.
  • Validation and downstream evaluation use fixed request orders because delayed inference is order-dependent.
  • 8.99 × 10−5 NLL is the largest difference among the three activation-scale policies at step 30,000.Each policy remains below the native Transformer Engine NVFP4 result of 2.325921.

B FP4 GEMM Simulation Details and Ablations

The FP4 GEMM simulation is fitted and tested with deliberately separating native witnesses, broad random probes, and full-model checks. The resulting model identifies 64-product groups and toward-zero accumulation as necessary for matching tested native behavior.

  • Seven real-model BF16-boundary witnesses compare custom and Transformer Engine paths with identical decoded operands, scales, and global scale products.
  • 0/7 decoded-operand Torch matmul outputs match the native witness set, making the cases diagnostic for accumulation behavior.
  • Only toward-zero combination reaches the native bin for the discriminating negative witness, while permuting K-blocks changes rounded outputs despite an unchanged mathematical dot product.
  • Groups of 64 are the only tested setting matching all 258 order tests, and the complete model matches 366 of 512 native FP32 values versus 6 for RNE accumulation.
  • The 1.2B full-model test matches every represented parameter-gradient value across one complete forward and backward pass.

C Statistical Consequences of the Probe-Matched Emulator

The emulator’s statistical probes isolate toward-zero contraction and final-grid tradeoffs, then test whether these choices alter optimization. Across compact and full-model controls, grid changes can be exactly invariant at the BF16 interface.

  • Toward-zero accumulation contracts magnitudes across Gaussian and heavy-tailed dot products, while heavier tails slightly increase the relative effect without changing its sign.
  • RTZ regression gradients preserve cosine similarity of at least 0.999999999996 while attenuating amplitude by roughly 1 ppm.
  • A final product-lattice grid maps encoded accumulators in [−1/2048, +1/2048] to zero, creating a real-value dead zone of total width α/1024.
  • 95.4% of genuine near-zero residuals are mapped to zero at the tested grid, lowering gain to approximately 0.44 and raising RMSE to 2.26 × 10−4.
  • All tested final grids produce identical BF16-interface values and optimizer trajectories through 500 updates in the tiny regression control.
  • After 100 full-model AdamW updates, one gradient element differs by two BF16 ULPs, while all final parameters and optimizer states remain byte-exact.
Loading 2609.02846v1…