PaperlayerView exact revision on arXiv

Source-linked AI summary

Training an MoE LLM in MXFP4: How We Closed the Gap to BF16

Han Wang, Zhitao Wang, Natalia Frumkin, Xuanwu Yin, Alireza Khodamoradi, Dong Li, Emad Barsoum

arXiv:2610.32229v2cs.AI

TL;DR

FP4 quantization can destabilize training, and ALTO composes MXFP4 quantization techniques with weight de-oscillation for GPT-OSS-20B. At global batch size 16, de-oscillation brings final validation loss to 3.3350 and reduces the steps-to-target overhead relative to BF16 to +5.0%.

  • Problem

    FP4 quantization error and weight oscillation challenge stable optimization during long-horizon training.

    Source: §1 · p. 1

  • Method

    ALTO combines hybrid block quantization, RHT, and stochastic rounding with late-training weight de-oscillation for MXFP4 training.

    Source: §1 · p. 2, §4 · p. 4

  • Results

    At GBS=16 and 16,128 steps, MXFP4 + RHT + SR + de-oscillation reaches validation loss 3.3350 (+0.007 vs. BF16) and cuts steps-to-target overhead from +20.0% to +5.0%.

    Source: §7 · p. 12

  • Takeaways & Limitations

    De-oscillation recovers most of the loss gap, supporting its use as a late-training stabilizer in the recommended recipe.

    Source: §7 · p. 12

    Experiments use simulated MXFP4 kernels on MI300X and do not report wall-clock speedup over BF16.

    Source: §6 · p. 11

Abstract

from arXiv · show

Training large language models (LLMs) in sub-8-bit arithmetic promises substantial gains in memory bandwidth and compute throughput, yet FP4 formats introduce quantization error severe enough to destabilize optimization if left unaddressed. We present ALTO, an open-source training recipe for GPT-OSS-20B, a 20-billion-parameter mixture-of-experts (MoE) model, on the MLPerf Small MoE Training benchmark under the MXFP4 microscaling format. The recipe combines established techniques from Abecassis, Felix et al. [arXiv:2509.25149] (hybrid 1D/2D block quantization, randomized Hadamard transforms, and stochastic rounding) with a simplified weight de-oscillation scheme adapted from TetraJet-v2 [arXiv:2510.27527]. On C4 validation at global batch size 16, MXFP4 + RHT + SR + de-oscillation reaches a validation loss of 3.3350 at 16,128 steps, only 0.007 above the BF16 baseline (3.3283), closing roughly 50% of the gap left by MXFP4 without de-oscillation (3.3418). Measured by steps to reach the MLPerf quality target (validation loss 3.34), de-oscillation cuts the FP4 convergence overhead relative to BF16 from +20.0% to +5.0%. We also report negative end-to-end results from differential gradient estimation, outlier clipping, and macro-block scaling, despite each showing operator-level SNR gains. All experiments use simulated MXFP4 kernels on AMD MI300X hardware and wall-clock speed has not been optimized.

1 Introduction

MXFP4 could reduce operand footprint and better exploit low-bit hardware, but its quantization error creates training-stability challenges. ALTO composes quantization and stabilization techniques for MXFP4 training of GPT-OSS-20B.

1 Introduction

  • MXFP4 uses 4-bit values with per-block scales, offering a theoretical 4× operand-footprint reduction relative to BF16.The format targets expert matrices that dominate compute and memory in sparse MoE models.Source: §1 · p. 1
  • Tiny master-weight updates can make quantized weights flip between adjacent bins, while outliers can shift shared block scales and trigger further flips.These weight oscillations can become more damaging during late, fine-grained convergence.Source: §1 · pp. 1–2
  • FP4 also amplifies sensitivity to tensor outliers, and straight-through estimation introduces gradient bias, compounding optimization challenges over long runs.Source: §1 · p. 2
  • ALTO composes hybrid 1D/2D block quantization, randomized Hadamard transforms, stochastic rounding, and weight de-oscillation for GPT-OSS-20B training.The recipe builds on methods from Abecassis, Felix et al. [1] and adapts de-oscillation from TetraJet-v2 [2].Source: §1 · p. 2

2 Related Work

The related work traces low-precision training from FP16 and FP8 toward FP4, where larger quantization errors require stabilization. ALTO combines and evaluates existing techniques for MoE training, alongside a simplified weight de-oscillation method.

Low-precision training

  • FP16 and FP8 established that reduced-precision formats can preserve convergence, while FP4 introduces substantially larger quantization error and requires additional stabilization.Source: §2 · p. 2

Microscaling and FP4 training

  • Prior FP4 training work explores hybrid block quantization, randomized Hadamard transforms, stochastic rounding, gradient estimation, outlier handling, and macro-block scaling.ALTO adopts and adapts several of these approaches, including methods from Abecassis, Felix et al. [1].Source: §2 · p. 2

Outlier mitigation and rotations

  • Prior methods mitigate outliers through mixed-precision isolation, error-compensated quantization, rotations, or clipping; ALTO applies rotation- and clipping-based approaches during training.The cited approaches include LLM.int8() [15], GPTQ [16], QuaRot [17], SpinQuant [18], and QuEST [19].Source: §2 · p. 2

Weight oscillation in quantized training

  • Weight oscillation describes quantized weights flipping between adjacent bins despite smooth full-precision trajectories; ALTO adapts TetraJet-v2’s [2] OsciReset for MX-format MoE training.Nagel, Markus et al. [20] previously analyzed and mitigated oscillation in quantization-aware training.Source: §2 · p. 2

Mixture-of-experts models

  • GPT-OSS-20B is the target of the MLPerf Small MoE training benchmark, and its expert matrices motivate FP4 training because they dominate compute and memory.Sparse MoE architectures scale model capacity while keeping per-token compute bounded [21] [22] [23].Source: §2 · p. 2

3 Problem Setting

The study evaluates FP4 training on the MLPerf Small MoE task using GPT-OSS-20B, a C4 dataset subset, and the benchmark’s validation-loss target. It compares operator-level numerical accuracy with end-to-end convergence.

3.1 Benchmark: MLPerf Small MoE Training

  • The benchmark evaluates from-scratch pretraining of GPT-OSS-20B, an open-source 20-billion-parameter MoE model, against a validation-loss target of 3.34.Training and validation use C4 splits, with preprocessed data of roughly 80 GB.Source: §3.1 · pp. 2–3

3.2 Training Protocol

  • The protocol follows the MLPerf task, dataset, target, and evaluation cadence, applying MXFP4/NVFP4 to linear and grouped-expert layers while excluding the lm_head and router gates.The hyperparameter table is identified as the GPT-OSS-20B pretraining configuration.Source: §3.2 · p. 3, Table 1 · p. 3

3.3 Evaluation Methodology

  • Operator tests compare MXFP4 outputs and gradients with a BF16 reference after forward and backward passes through an MSE loss.Synthetic inputs, weights, and targets include sparse, roughly 100×-standard-deviation outlier perturbations at 0.5% probability.Source: §3.3 · p. 3
  • Operator-level signal-to-noise ratio is reported in decibels for forward output, input gradient, and weight gradient; end-to-end convergence is measured by steps to first reach validation loss ≤3.34.Validation checkpoints occur every 768 steps at global batch size 16, so the earliest recorded crossing is reported.Source: §3.3 · p. 3

4 Method

ALTO combines hybrid block quantization, randomized Hadamard transforms, and stochastic rounding, with weight de-oscillation applied late in training. The hybrid layout addresses differing weight access patterns in forward and backward passes while preserving an activation path compatible with RHT.

4 Method

  • The recommended stack combines hybrid block quantization, RHT, and stochastic rounding, then applies weight de-oscillation during late training.The first three techniques are adopted from Abecassis, Felix et al. [1], while de-oscillation adapts TetraJet-v2 [2].Source: §4 · p. 4
  • MXFP4 uses shared block scales and E2M1 values; where native MXFP4 GEMMs are unavailable, simulated quantize-dequantize operations reproduce numerical error without bandwidth savings.The layer’s three GEMMs produce forward output, activation gradient, and weight gradient; gradients use straight-through estimation unless modulated by DGE.Source: §4 · p. 4

4.1 Hybrid 1D/2D Block Quantization

  • A 1D weight layout requires quantization along different axes for forward and backward passes, creating extra overhead and inconsistent weight views.Source: §4.1 · p. 4
  • The hybrid layout keeps activations in canonical 1D blocks compatible with RHT and uses 2D weight blocks that serve both forward and backward passes.The method extends the MX specification for weights while retaining its activation layout.Source: §4.1 · p. 5

4.2 Randomized Hadamard Transform

  • RHT transforms activations and output gradients on the weight-gradient path to redistribute outlier energy before quantization.The transform is orthogonal up to scaling and is applied on the 1D activation path, making it incompatible with 2D activation blocking.Source: §4.2 · p. 5

4.3 Stochastic Rounding on Gradients

The recipe uses stochastic rounding for gradient quantization and adds optimizer-level weight de-oscillation to address quantized-weight oscillation. De-oscillation flags elements by their quantized-to-full-precision travel ratio and snaps flagged weights to bin centers.

4.3 Stochastic Rounding on Gradients

  • Stochastic rounding makes gradient quantization an unbiased estimator of the pre-quantization gradient in expectation, while forward-pass weight and activation quantization remains deterministic.Elements round up or down with probabilities proportional to their distances from adjacent representable values.Source: §4.3 · p. 5

4.4 Weight De-Oscillation

  • De-oscillation flags an element when DistRatio = dQ/dw exceeds τ, then snaps flagged weights to their current bin centers at each window’s end.The default period is P = 200 steps and the default threshold is τ = 4.0.Source: §4.4 · pp. 5–6
  • ALTO implements de-oscillation as a model-agnostic AdamW post-step hook, activating after a 2,000-step warmup at GBS=16 with τ = 4.0.These settings activate earlier and use a lower threshold than TetraJet-v2’s 8,000-step start and τ = 16; both are tuned to ALTO’s MXFP4 training horizon.Source: §4.4 · p. 6

4.5 Differential Gradient Estimation (DGE)

ALTO uses a continuous piecewise power-law surrogate for differential gradient estimation, while also supporting static and dynamic outlier clipping. Dynamic clipping requires co-use with randomized Hadamard transforms in this implementation.

4.5 Differential Gradient Estimation (DGE)

  • ALTO’s corrected DGE uses a piecewise power-law surrogate that is continuous between segments, unlike the cited estimator’s discontinuities at bin boundaries.For MXFP4, bin widths and midpoints come from the E2M1 representable grid; f′(x) is clamped to [0, 3] before application.Source: §4.5 · p. 6

4.6 Outlier Clipping

  • Outlier clipping caps large-magnitude elements before quantization to reclaim precision for the majority of elements.A few outliers can inflate the per-block scale, compressing or saturating the remaining values.Source: §4.6 · p. 7
  • Static clipping scales inputs by 3/4 before quantization and by 4/3 after dequantization, with a 16/9 compensating factor on the weight-gradient path.Dynamic clipping estimates a per-block threshold from block standard deviation and requires RHT in this implementation.Source: §4.6 · p. 7

4.7 Macro-Block Scaling

Macro-block scaling adds a coarse scale before standard MXFP4 block quantization, while low-rank outlier compensation adds a branch to preserve information lost to FP4 saturation. The latter is applied only to linear layers, not MoEs.

4.7 Macro-Block Scaling

  • Macro-block scaling applies a coarse macro-scale before standard 32-element MXFP4 quantization to better preserve outliers.ALTO uses 128 × 128 blocks for weights and 1 × 128 blocks for activations, mapping each macro-block maximum to E2M1 magnitude 6.0.Source: §4.7 · p. 7
  • The macro-scale is stored as a shared mantissa using only the upper 8 bits of the FP32 scale, then reconstructed for prescaling and descaling.The scheme addresses the coarse granularity of one UE8M0 scale shared across each 32-element MX block.Source: §4.7 · p. 7

4.8 Low-Rank Outlier Compensation

  • Low-rank outlier compensation decomposes a linear layer into a full MXFP4 GEMM plus a rank-r branch intended to capture information lost to FP4 saturation.Each compensated layer adds MXFP4 factors U and V, an FP32 parameter σ, and six additional low-rank GEMMs per forward-backward pass; MoEs are not considered.Source: §4.8 · p. 7

4.9 M+Adam: Multiplicative Update for Adam

The described Adam update forms moment estimates from quantized gradients while retaining optimizer state and updates in higher precision. M+Adam adds a multiplicative branch alongside the standard additive update, at the cost of extra per-weight state and computation.

4.9 M+Adam: Multiplicative Update for Adam

  • In low-precision training, Adam forms its moment estimates from quantized gradients while retaining optimizer state and updates in higher precision.The first- and second-moment estimates use decay rates β1 and β2, and the additive update is identified as standard Adam.Source: §4.9 · p. 8
  • M+Adam augments the standard additive Adam update with a multiplicative branch based on each nonzero weight’s mantissa-like and exponent-like coordinates.The supplied passage describes an additional second-moment accumulator normalizing the exponent-coordinate gradient to produce a relative update.Source: §4.9 · p. 8
  • The multiplicative update scales with current weight magnitude to help updates cross low-precision quantization intervals, but requires additional per-weight optimizer state and computation.The additive update remains necessary near zero and for sign changes.Source: §4.9 · p. 8

5 Results

At GBS=16, the recommended MXFP4 recipe with de-oscillation approaches BF16 validation loss and reaches the 3.34 target with a smaller step overhead than other FP4 recipes. Operator-level SNR gains did not consistently predict end-to-end validation loss, while de-oscillation provided a late-training benefit.

5.1 Operator-Level Numerical Accuracy

  • SR lowers input-gradient SNR but removes systematic rounding bias, and that bias reduction outweighs the local SNR penalty in end-to-end training.RHT leaves the forward pass unchanged and modestly improves weight-gradient SNR.Source: §5.1 · pp. 8–9
  • MBS gives the largest operator-level SNR gain, while low-rank compensation modestly raises forward and input-gradient SNR but does not improve end-to-end convergence at r = 32.The compensated layer's weight gradient is not comparable to the plain BF16 reference; its forward output and input gradient remain close.Source: §5.1 · pp. 8–9

5.3 Technique Evaluation and Recipe Selection

  • De-oscillation delivers the largest ablation improvement (−0.008), while DGE and dynamic clipping are strongly detrimental (+0.201 and +0.181, respectively).Static clipping is neutral (∆= 0.000), and low-rank compensation at r = 32 is slightly harmful (+0.007).Source: §5.2 · pp. 9–10, Table 5 · p. 10

5.2 End-to-End Training

  • Within 0.007 of BF16 validation loss at 16k steps, the recommended MXFP4 recipe reaches the 3.34 target with +5.0% step overhead, versus +20.0% for other FP4 recipes.The recommended configuration is 1d2d + RHT + SR + de-oscillation; Table 4 reports steps to the target.Source: §5.2 · pp. 9–10, Table 4 · p. 10, Figure 4 · p. 10

5.3 Technique Evaluation and Recipe Selection

  • Operator-level SNR is an imperfect proxy for end-to-end quality: MBS and static clipping raise SNR but only marginally affect validation loss, while DGE and dynamic clipping are harmful.De-oscillation provides the most reliable late-training benefit among the techniques evaluated.Source: §5.3 · p. 11
  • The preferred stack is 1d2d + RHT + SR, paired with de-oscillation for closed-division settings or M+Adam for open-division settings.The benchmark's closed division requires exactly AdamW, whereas open-division optimizer choice is unrestricted.Source: §5.3 · p. 11

6 Limitations

The evaluation measures algorithmic robustness, not production-stack speed: simulated MXFP4 does not capture hardware savings, and the training path is unoptimized. Results are also limited to one model family, data subset, and accelerator setting.

6 Limitations

  • MXFP4 results model numerical error, not memory-bandwidth or compute savings, because MI300X execution uses simulated quantize/dequantize round trips.Separate, unfused quantization, RHT, and GEMM launches also prevent a wall-clock speedup comparison with BF16.Source: §6 · p. 11
  • The training path is not throughput-optimized; production speed would require operator fusion, persistent MXFP4 buffers, and removal of redundant quantization/dequantization.Low-rank outlier compensation adds per-layer GEMMs each step, while de-oscillation mainly adds optimizer-state memory that fits available headroom.Source: §6 · pp. 11–12
  • Evaluation covers only GPT-OSS-20B MoE on a C4 subset using AMD MI300 accelerators; generalization to dense models, other FP4 formats, larger batches, and longer runs remains unestablished.The authors also call for further study of low-rank compensation across ranks and on grouped GEMM.Source: §6 · p. 12

7 Conclusion

The recommended MXFP4 recipe combines 1D/2D block quantization, RHT, and SR with weight de-oscillation. At GBS=16, de-oscillation narrows the BF16 loss gap and reduces steps-to-target overhead, while ablations show that operator SNR gains do not guarantee better end-to-end loss.

7 Conclusion

  • At GBS=16 and 16,128 steps, de-oscillation reaches validation loss 3.3350 (+0.007 vs. BF16), and reaching the 3.34 quality target costs +5.0% more steps than BF16.The baseline MXFP4 stack reaches 3.3418 versus 3.3283 for BF16 (+0.014); de-oscillation closes half the remaining loss gap with modest optimizer-state memory and no extra GEMMs.Source: §7 · p. 12
  • Higher operator SNR does not ensure better end-to-end training: MBS and static clipping do not improve validation loss at GBS=16, while DGE and dynamic clipping are harmful at scale.Low-rank compensation at r = 32 is neutral to slightly harmful at GBS=16 and trails de-oscillation (∆= −0.008) in accuracy and compute efficiency.Source: §7 · p. 12
  • Future work targets fused quantization and transforms, FP4 dispatch in expert parallelism, and an audited MLPerf Training submission on native FP4 hardware.Source: §7 · p. 12

Related papers

Loading 2610.32229v2…