🤖 machine learning

Pretraining large language models with MXFP4

This paper identifies that quantizing weight gradients (Wgrad) is the primary cause of instability in full-pipeline FP4 training of large language models and demonstrates that deterministic Hadamard rotations, rather than stochastic rounding, are essential to restore convergence stability.

Original authors: Musa Cim, Poovaiah Palangappa, Miro Hodak, Ravi Dwivedula, Meena Arunachalam, Mahmut Taylan Kandemir

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Musa Cim, Poovaiah Palangappa, Miro Hodak, Ravi Dwivedula, Meena Arunachalam, Mahmut Taylan Kandemir

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, brilliant student (a Large Language Model) to write stories. Usually, to keep this student learning effectively, you give them notes written in high-quality, detailed ink (called FP8 precision). This works well, but the notes are heavy, taking up a lot of space in the student's backpack and slowing down how fast they can read them.

Researchers wanted to see if they could switch to ultra-lightweight notes written in just 4 bits of information (MXFP4). These notes are tiny and fast, promising to make the student learn much quicker. However, there's a catch: when they tried using these tiny notes for the entire learning process, the student often got confused, started making wild guesses, and failed to learn anything (the training "diverged").

The paper asks: Why does this happen, and how do we fix it?

The Experiment: A Three-Step Test

The researchers set up a controlled experiment to figure out exactly where the problem lies. They broke the learning process into three stages and slowly swapped the heavy notes for the light notes, one stage at a time:

Forward Pass (Fprop): The student reads the question.
Activation Gradients (Dgrad): The student figures out how much they misunderstood the question.
Weight Gradients (Wgrad): The student updates their brain (the "weights") to remember the lesson for next time.

The Findings:

Stage 1 & 2: When they used light notes just for reading and figuring out misunderstandings, the student did fine. They learned almost as well as with the heavy notes, just needing a few extra practice questions.
Stage 3 (The Problem): The moment they switched to light notes for updating the brain (Wgrad), the student started to spiral out of control. The researchers found that this specific step is the "weak link." It's like trying to write a complex math correction on a tiny, crumpled piece of paper; the details get lost, and the student learns the wrong lesson.

The Failed Fixes: Adding Randomness

When the student started failing, the researchers tried a common trick: Stochasticity (adding randomness).

The Analogy: Imagine the student is confused, so you tell them, "Just guess randomly!" or "Spin a wheel to decide your answer."
The Result: This didn't work. In fact, adding randomness to the tiny notes made the errors worse. The "noise" of the random guesses amplified the tiny errors already present in the light notes, causing the training to crash completely.

The Winning Fix: The Deterministic Shuffle

The researchers then tried a different approach: Deterministic Hadamard Rotations.

The Analogy: Instead of guessing randomly, imagine the student has a special, rigid ruler that perfectly rearranges the tiny notes before they are written down. This ruler doesn't add randomness; it simply organizes the information in a specific, predictable way so that the tiny notes don't get crumpled or lost.
The Result: This worked perfectly. By using this specific, predictable "rearranging" method, the student could use the tiny, fast notes for the entire process (reading, analyzing, and updating) without getting confused.

The Bottom Line

The paper concludes with two main takeaways:

The Bottleneck: The main reason 4-bit training fails isn't because the notes are too small generally; it's specifically because the brain-updating step (Wgrad) is too sensitive to the tiny errors in those notes.
The Solution: To make 4-bit training work, you don't need to add more randomness. You need structure. By using a specific, predictable mathematical shuffle (Hadamard rotation), you can stabilize the process.

The Payoff:
When they fixed the problem with this "rearranging ruler," the system became 20% faster at each step of learning. Because the student learned stably without needing extra practice questions, the whole process finished about 10% faster from start to finish compared to the old, heavy-ink method.

In short: You can use the super-fast, tiny notes, but only if you organize them perfectly before writing them down. If you just throw in some randomness, it breaks everything.

Technical Summary: Pretraining Large Language Models with MXFP4

Problem Statement

The rapid scaling of Large Language Models (LLMs) has intensified the demand for hardware-efficient training algorithms. While FP8 has emerged as a viable low-precision alternative, training with 4-bit formats (such as MXFP4) remains a significant challenge. A primary difficulty is that activation and gradient outliers amplify quantization errors, potentially destabilizing optimization. Early studies indicate that coarsely quantizing activation outliers can dominate the quantization range, severely degrading floating-point quality. Despite efforts to mitigate noise via mixed precision or staged schedules, a fundamental question remains: Why does full-pipeline FP4 training often diverge even when forward activations and activation gradients appear stable? Furthermore, it is unclear which specific components of the training pipeline (Forward propagation, Activation gradients, or Weight gradients) are most responsible for convergence degradation and what stabilization mechanisms are effective.

Methodology

The authors conducted a controlled study on the pretraining of Llama 3.1–8B using the C4 dataset, adhering to the MLPerf LLM pretraining setup. The study was executed on AMD Instinct MI355X GPUs utilizing native FP4 tensor-core support, distinguishing this work from prior research that relied on software emulation.

The core experimental design involved a progressive "ladder" of MXFP4 enablement:

Baseline: Full FP8 training.
Stage-wise Quantization: MXFP4 was progressively enabled across three components while holding others fixed at FP8:
- Forward propagation (Fprop) only.
- Fprop + Activation gradients (Dgrad).
- Fprop + Dgrad + Weight gradients (Wgrad).
Stabilization Evaluation: Under the full MXFP4 pipeline (Fprop + Dgrad + Wgrad), the authors tested various interventions to restore stability:
- Stochastic Rounding: To introduce randomness.
- Randomized Hadamard Rotations: To decorrelate outliers.
- Deterministic Hadamard Rotations: Specifically using $16 \times 16$ (H16) and $32 \times 32$ (H32) matrices.

Metrics: Convergence was defined as reaching a validation perplexity of 3.3. The primary metric for efficiency was token overhead (the relative increase in tokens required to reach the target compared to the FP8 baseline). Secondary metrics included training-step throughput and end-to-end speedup.

Key Contributions and Findings

1. Identification of the Instability Source

The study isolates Weight Gradient (Wgrad) quantization as the primary driver of convergence degradation.

Enabling MXFP4 for Fprop only incurred an 8–9% token overhead.
Enabling Fprop + Dgrad resulted in a similar 10–11% overhead.
Enabling the full pipeline (Fprop + Dgrad + Wgrad) without stabilization caused the token overhead to jump to 26–27% or resulted in complete divergence ("Does Not Converge").
Conclusion: The instability is not inherent to Fprop or Dgrad quantization but is dominated by the quantization of Wgrad.

2. Evaluation of Stabilization Strategies

The authors tested whether adding stochasticity or structured rotations could stabilize the full pipeline:

Stochastic Rounding: Failed to stabilize the full pipeline; the run diverged.
Randomized Hadamard Rotations: Similarly failed to converge when applied to the full pipeline.
Deterministic Hadamard Rotations: Successfully restored stable optimization. Applying a deterministic Hadamard transform (specifically H16) to the full MXFP4 pipeline reduced the token overhead back to 8–9%, closely tracking the FP8 baseline.
Insight: Stochastic interventions appear to amplify effective quantization error rather than mitigating outlier effects. The instability is driven by structured micro-scaling errors along sensitive gradient paths, which deterministic rotations can correct.

3. Efficiency and Throughput

With the stability issue resolved via Deterministic Hadamard (H16), the authors demonstrated the practical efficiency gains of FP4:

Throughput: The MXFP4 + H16 configuration achieved a 20% higher training-step throughput compared to the FP8 baseline on MI355X GPUs.
End-to-End Speedup: Because the token overhead was minimized (8–9%), the throughput advantage translated into a 9–10% end-to-end speedup.
Key Insight: FP4 speedups are "stability-gated." Without controlling Wgrad instability, the increased token-to-converge requirement negates the raw compute benefits of lower precision.

Significance and Claims

The paper claims to provide the first demonstration of full MXFP4 pretraining on MI355X-class hardware without software emulation. Its significance lies in shifting the understanding of FP4 training dynamics:

Mechanism of Failure: Instability in FP4 pretraining is not due to insufficient stochasticity but rather structured micro-scaling errors in sensitive gradient paths (specifically Wgrad).
Practical Recipe: A deterministic Hadamard rotation is a necessary and effective stabilizer for full-pipeline MXFP4 training, enabling the realization of hardware efficiency gains.
Contextual Limitation: The authors modestly note that this stabilization recipe is setting-dependent. While effective for Llama 3.1–8B on the C4 dataset, the behavior may vary across different models, datasets, or fine-tuning methods, and should not be assumed universal without further evidence.

In summary, the work establishes that replacing FP8 with MXFP4 is feasible for large-scale pretraining, provided that Wgrad quantization errors are controlled via deterministic structural interventions rather than stochastic noise injection.