🤖 machine learning

Max-Window Scale Estimation for Near-Lossless HiF8 W8A8 Quantization-Aware Training

This paper identifies and resolves two distinct failure modes in HiF8 W8A8 quantization-aware training—amax saturation and catastrophic forgetting—by introducing a max-window scale estimation strategy and a BF16 warmup schedule, thereby achieving near-lossless performance on the OpenPangu-Embedded-1B model.

Original authors: Yingying Cheng, Jinquan Shi, Li Zhou, Zhiyang He, Zhaoyi Sun, Fan Zhang, Jie Sun

Published 2026-05-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yingying Cheng, Jinquan Shi, Li Zhou, Zhiyang He, Zhaoyi Sun, Fan Zhang, Jie Sun

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 have a brilliant, well-read librarian (the AI model) who has spent years memorizing a vast library of books. Now, you want to shrink this massive library down so it fits into a tiny, portable backpack (a low-power device) without losing any of the stories or facts inside. This process is called Quantization.

The paper describes a specific experiment where the researchers tried to shrink a 1-billion-parameter AI model using a new, compact format called HiF8 (High-precision Float 8). Think of HiF8 as a "smart compression" that keeps the most important details sharp while compressing the rest.

However, the researchers discovered that simply shrinking the library wasn't enough. They found two hidden "traps" that could ruin the model's memory, even if the training process looked perfect on paper.

Here is the breakdown of their journey, explained simply:

The Two Hidden Traps

1. The "Blind Spot" Trap (Amax Saturation)
Imagine the librarian is trying to organize books on a shelf, but they are using a ruler that is slightly too short.

The Problem: The AI needs to know the "loudest" or "largest" number it sees to set the scale for compression. The researchers used a method called Delayed Tensor Scaling (DTS). This is like the librarian looking at the previous book to guess the size of the current book.
The Failure: If the current book is suddenly huge (a "spike" in data), the librarian's guess (based on the previous book) is too small. The book gets "clipped" or cut off at the edge of the shelf.
The Sneakiness: The AI's internal "scorecard" (training loss) didn't show this happening. It was like the librarian saying, "I'm organizing fine!" while secretly tearing pages out of the books. The damage only showed up later when they tested the librarian's knowledge on specific quizzes (like ARC or MMLU).

2. The "Over-Eager Student" Trap (Catastrophic Forgetting)
Imagine the librarian is so eager to learn new stories from a new set of books that they start erasing the old classics from their memory to make room.

The Problem: The researchers were teaching the model with a "learning rate" (how fast it learns) that was too aggressive.
The Failure: The model was learning the new data so fast that it completely forgot the common sense and facts it learned during its original training.
The Sneakiness: This happened even without the compression. If they just trained the model normally at this high speed, it still forgot everything. But because they were also compressing the model, they blamed the compression for the failure, not the speed.

The Solution: A Two-Part Fix

The researchers ran eight different experiments to figure out how to fix these traps. They found that fixing just one wasn't enough; they had to fix both simultaneously.

Fix #1: The "Safety Net" (Max-Window Strategy)
To stop the "Blind Spot," they changed how the librarian guesses the book sizes.

Instead of looking at just the previous book, they told the librarian to look at the largest book seen in the last 64 books.
This is a "conservative" approach. It makes the shelf slightly bigger than strictly necessary (a little over-conservative), but it guarantees that no book ever gets cut off. This small "extra space" actually helped the model stay stable, acting like a gentle regularizer.

Fix #2: The "Cool-Down" (Warmup and Slower Learning)
To stop the "Over-Eager Student," they changed the training schedule.

The Warmup: For the first 500 steps, they didn't compress the model at all. They let it get comfortable with the new data in its full, high-quality format (BF16). This let the model settle into a stable state before the "backpack" compression was applied.
The Slow Down: They drastically slowed down the learning speed (lowering the learning rate). This prevented the model from frantically overwriting its old, valuable knowledge with new data.

The Result

When they combined these two fixes, the result was near-lossless.

The compressed model (HiF8) performed almost exactly as well as the uncompressed, full-size model.
The drop in performance on difficult quizzes was tiny (less than 0.5% in most cases).
Crucially, they proved that if you only used the "Safety Net" but kept the "Over-Eager" speed, the model still failed. And if you slowed down the speed but kept the "Blind Spot" ruler, it also failed. Both fixes were necessary.

What Didn't Work (And Why)

Looking at just the last book: Caused books to get cut off.
Smoothing the guess: Tried to average past books, but failed when the data changed unpredictably.
Checking the current book instantly: This required a second look at the data, which was too slow and caused the shelf size to jump around too much, confusing the model.
High-speed training: Even without compression, this made the model forget everything.

The Bottom Line

The paper teaches us that when shrinking a smart AI model, you can't just look at the "training score" to see if it's working. You have to be careful about how you measure the data (to avoid clipping) and how fast you teach it (to avoid forgetting). By using a "safety net" for the data size and a "slow, steady" learning pace, you can fit a giant brain into a tiny backpack without losing its mind.

Technical Summary: Max-Window Scale Estimation for Near-Lossless HiF8 W8A8 Quantization-Aware Training

Problem Statement

The deployment of Large Language Models (LLMs) in resource-constrained environments necessitates aggressive quantization, particularly using 8-bit floating-point (FP8) formats like HiF8 (High-precision Float 8). While Quantization-Aware Training (QAT) aims to adapt pretrained models to low-bit representations, it introduces subtle failure modes that are invisible to standard training metrics like cross-entropy loss.

Specifically, applying QAT to pretrained models involves a complex trade-off: preserving pre-trained knowledge while adapting internal representations to quantization noise. The authors identify that standard training loss often fails to signal degradation, leading to two distinct, orthogonal failure modes in HiF8 W8A8 QAT:

$A_{max}$ Saturation: Caused by the temporal lag in Delayed Tensor Scaling (DTS), where the estimated scale factor is based on historical data. If the current step's true activation maximum exceeds the predicted scale, values are clipped, corrupting forward activations while the Straight-Through Estimator (STE) incorrectly passes gradients as if no clipping occurred.
Catastrophic Forgetting: Caused by an overly aggressive learning rate during continued pretraining, which overwrites pre-trained commonsense knowledge independently of the quantization scheme.

Methodology

The study focuses on the OpenPangu-Embedded-1B model (1 billion parameters) using the HiF8 W8A8 format. The methodology involves a systematic investigation through eight controlled experiments to disentangle the causes of performance degradation.

Quantization Configuration

Format: HiF8 uses tiered mantissa precision based on value magnitude (3 bits for $|x| \le 8$ , 2 bits for $|x| \le 128$ , 1 bit for larger values).
Module Selection: To balance efficiency and stability, only MLP layers (gate, up, and down projections) in the middle 21 Transformer blocks are quantized. Attention projections, embeddings, layer norms, and the boundary blocks (first 3 and last 2) remain in BF16 to preserve outlier handling and output precision.
Training Protocol: Experiments utilize the FineWeb dataset with a global batch size of 1024 and sequence length of 1024 over 10,000 steps.

Experimental Variables

The core of the investigation compares different $A_{max}$ estimation algorithms and learning rates:

$A_{max}$ Algorithms:
- Most-Recent (DTS): Uses the previous step's maximum.
- Exponential Smoothing (DTS): Uses a running average.
- Max-Window (DTS): Uses the maximum value over a 64-step history window.
- Current Tensor Scaling (CTS): Computes the maximum from the current step (requires a second forward pass).
Learning Rates: Experiments tested $10^{-4}$ (initially deemed safe for BF16 baselines) and $10^{-5}$ .
Warmup: A 500-step BF16 warmup phase was introduced in later experiments to stabilize weights and populate the $A_{max}$ history before activating quantization.

Key Contributions

The paper makes four primary contributions based on the disentanglement of failure modes:

Identification of Orthogonal Failure Modes: The authors demonstrate that benchmark degradation in HiF8 QAT stems from two independent sources: $A_{max}$ saturation (a quantization-specific issue) and catastrophic forgetting (a training-dynamics issue). Crucially, neither is detectable via training loss alone; both manifest only in downstream evaluation.
Conservative Max-Window Strategy: To address $A_{max}$ saturation, the paper proposes replacing standard algorithms with a max-algorithm over a 64-step history window. This provides a hard guarantee against scale-induced saturation for any step within the observation window. The slight over-conservatism acts as beneficial regularization.
Two-Phase Training Protocol: To mitigate catastrophic forgetting, the authors introduce a protocol involving a 500-step BF16 warmup followed by QAT at a reduced learning rate of $10^{-5}$ . The warmup ensures the model reaches a stable weight regime and populates the $A_{max}$ history with representative values, while the lower learning rate prevents the overwriting of pre-trained knowledge.
Comprehensive Failure Analysis: The study documents why previous approaches failed:
- Most-Recent suffers from 1-step delay saturation.
- Exponential Smoothing fails under non-monotone activation dynamics.
- CTS introduces per-step fluctuations that corrupt stable representations.
- High Learning Rates ( $10^{-4}$ ) cause significant forgetting regardless of quantization.

Results

The final configuration (Max-Window algorithm + 500-step warmup + $lr=10^{-5}$ ) achieves near-lossless quantization compared to a matched BF16 baseline trained at the same learning rate:

MMLU: 0.43% drop (43.17% vs. 43.36% baseline).
HellaSwag: 0.58% drop.
ARC-Challenge: 0.22% drop.
Training Loss: The Average Percentage Error (APE) between the HiF8 QAT and BF16 baseline is only 0.11% over 10,000 steps.

Notably, experiments showed that neither fix alone was sufficient. Using the Max algorithm at $lr=10^{-4}$ still resulted in ~3% drops on ARC tasks due to forgetting. Conversely, using $lr=10^{-5}$ with a standard most-recent algorithm would still suffer from saturation.

Significance and Claims

The paper claims that successful HiF8 QAT for pretrained models requires addressing the interaction between scale stability and learning dynamics simultaneously. The significance of the work lies in:

Diagnosing the "Invisible" Failure: Highlighting that standard training loss is an insufficient metric for QAT quality, as it masks saturation and forgetting that only appear in downstream benchmarks.
Providing a Robust Solution: Demonstrating that a conservative $A_{max}$ estimation strategy combined with a lower learning rate and warmup can achieve near-lossless performance (sub-1% drop on major benchmarks) for 1B parameter models.
Clarifying Design Axes: Establishing that the choice of $A_{max}$ algorithm and learning rate are critical, interdependent design axes, rather than isolated hyperparameters.

The authors conclude that while the current configuration is effective for 1B models, future work may explore layer-wise learning rates, adaptive warmups, and increasing the number of high-precision layers to further close remaining gaps.