Making Quantized Models More Accurate

We recently incorporated quantization-aware distillation (QAD) into our training platform. This technique allows us to recover a significant amount of accuracy for quantized models, often reducing the held-out KL of NVFP4 models against their BF16 base by upwards of 40%.

Surprisingly, large KL reductions did not always correspond to recovered task accuracy, and we documented numerous instances of large reductions in KL with little task accuracy recovery, as well as small reductions in KL with much more. This opens the door to a new research area on what other metrics might be better proxies for preserving evaluation accuracy following quantization.

In this blog post, we’ll share some of the experiments we conducted on our way to better quantized models, as well as early results and open questions.

Why Quantize?

Quantization replaces high-precision numbers with a smaller set of representable values. For instance, you might map 16-bit floating point numbers to 4-bit floating point numbers. This mapping happens via a multiplicative factor called a scale and an optional affine shift. The grouping of numbers assigned to a single scale is the granularity or block-size. Different formats quantize with one scale per tensor, per row, or even per small segment within a row. [1]⌝

Quantization makes models cheaper in two ways:

1. It reduces the amount of data we need to move around the chip, reducing memory bandwidth pressure.
2. It reduces the cost of matrix multiplications themselves, increasing arithmetic throughput through specialized hardware acceleration.

Quantization can be weights-only or both weights and activations. Quantizing weights typically achieves only the first benefit, as weights are de-quantized to perform matrix multiplications with higher precision activations. Since decode is mostly memory-bandwidth bound, this still reaps significant gains. Quantizing weights and activations can also accelerate compute-bound work, particularly prefill and larger-batch decode. We quote quantization using WA format, so W8A16 means 8-bit weights and 16-bit activations. [2]⌝

Quantization is a beautiful example of software-hardware co-design:

• Ampere SM80 implements native BF16 and FP16 tensor cores. [3]⌝
• Hopper SM90 implements native FP8 tensor cores. [4]⌝
• Blackwell SM100 implements native block-scaled FP4 and FP8 tensor cores. [5]⌝

On Blackwell, the tcgen05 family includes block-scaled matrix multiply-accumulate instructions. This allows the GPU to consume packed values and block scales together, performing a GEMM without needing to materialize a full BF16 tensor before the multiply. [6]⌝

While quantization provides big speedups and reduces memory bottlenecks, there is an obvious drawback: quantization can degrade task performance. The extent to which this happens depends heavily on the quantization format, which components you’ve quantized, how the model was trained, and what task you’re evaluating the model on. [7]⌝

At Applied Compute, we want to train the best models possible for specific tasks, but we also want to deploy them with the most compelling economics. Thus, we’re very interested in minimizing the accuracy degradation that comes with quantization.

What to Quantize?

One way to preserve accuracy is to keep some components of the model in higher precision while quantizing others more aggressively. You can apply blanket rules, such as keeping all attention components in BF16, or make specific decisions, such as preserving the MLP gate projection in layer 42. This is called mixed-precision quantization.

The challenge is deciding where the additional precision matters. This reduces to a constrained optimization problem of maximizing accuracy given an average bit width budget.

We explored this with a mixed-precision search on a Qwen 3.8 27B, allowing each eligible component to use NVFP4, FP8, or BF16. There are two ways to estimate which components need higher precision:

1. KL against BF16. Quantize one component or group of components while leaving the rest of the model in BF16. Then, measure the change in the model’s output distribution. Repeat this for every component and candidate format. This directly measures the effect on output KL, but requires full model forward passes for every choice. Our setup had 976 group–format combinations across 128 scoring examples, requiring roughly 125,000 full model forwards.
2. Gradient sensitivity. Run a forward and backward pass to measure how sensitive the loss is to each component’s output. Then replay each component locally under the candidate quantization formats. Score its quantization error, weighted by the squared loss gradient:

Here, $y_i$ is the component’s BF16 output, $\widehat y_{i,f}$ is its output under format $f$, and $\mathcal L$ is the teacher-forced cross-entropy loss. The sum runs over examples and output elements. Errors receive a larger penalty where the loss is more sensitive. This avoids rerunning the entire model for every candidate, but ignores interactions between output errors. [8]⌝

We’ve experimented much more with the gradient sensitivity approach, as it completes in an order of magnitude less time. In one experiment, we calibrated scales on 400 full-context examples from the BF16 policy and estimated sensitivity on 128 examples using 4,096-token prefixes, computing loss over the final 512 predictable tokens in each prefix. We then reused these scores to find allocations that minimized total gradient sensitivity under different bit budgets.

These searches produce clear patterns. Earlier layers tend to retain higher precision, and attention is generally preserved before MLPs. Some components, including the linear-attention A/B and convolution paths, were excluded from the search and kept in BF16 throughout. At a 10-bit budget, one mixed checkpoint occupied 36.46 GB, compared with 55.56 GB for BF16: a 34.4% reduction. Measured performance was essentially unchanged on the full 80-task evaluation.

But they don’t consistently beat simpler choices. In many instances, we found that uniform FP8 beat a mixed allocation at the same bit budget, despite scoring worse on the gradient proxy. The allocation that minimized the proxy was not necessarily the one that best matched the original model.

We’ve found that blanket rules for which components should be high precision work nearly as well and with very little computational cost compared to running a proper mixed-precision optimization sweep. In general, MoE routers, shared experts, and attention components should be retained in high precision. Quantizing routed experts delivers much of the performance gains while recovering significant lost accuracy from full-component quantization. This aligns with findings discussed previously in our Kimi K3 technical write-up⌝, where we discovered that expert-only quantization checkpoints looked quite similar to unquantized checkpoints in KL terms.

In one example, we tested this on Qwen 3.8 27B using an internal multi-turn reasoning dataset. The model observes an environment state and submits actions through a tool. Across W8A16, W8A8, W4A16, and W4A4 recipes, we compared quantizing both attention and MLP projections with quantizing only the MLPs.

Attention projections accounted for approximately 17% of the model’s parameters, but keeping them in BF16 reduced KL against the original model by 49-54% across all four recipes on 180 held-out traces.

This provides a simple starting point: preserve attention and quantize the MLPs. Experiment with KV cache in FP8 before moving to NVFP4.

We next investigated whether we could train the model to recover its lost accuracy.

Quantization-Aware Distillation

In quantization-aware distillation (QAD), the original unquantized model serves as a frozen teacher. We construct the student by quantizing the teacher but otherwise initializing it from the same weights. We then train the quantized student’s outputs to match the teacher distribution. Prior to this, neither the student nor teacher had been trained in a quantized setting. The goal is to recover behavior that would otherwise be lost when we quantize the model. [9]⌝

A common implementation uses quantize-dequantize, or QDQ, also known as fake quantization. This technique rounds weights and activations to their low-precision representations and then converts them back to high precision before the matrix multiplication. This introduces quantization error but allows you to train using existing high-precision kernels. [10]⌝ The backward pass uses a straight-through estimator (STE), implemented using stop-gradient, through rounding. [11]⌝ We employ both native quantized training as well as QDQ, but we opted for QDQ here as native quantized training⌝ introduces stability problems that must be controlled for.

Training-inference mismatch

Training with QAT or QAD generally has higher training-inference mismatch compared to ordinary BF16 training (even with staleness of zero). One way to measure training-inference mismatch is to compute the forward KL between the same precision policy in the training engine versus the rollout engine. This tells you how strongly outputs change when you vary the engine and nothing else. Similarly, you can hold the engine constant and vary only the quantization, measuring the KL introduced by quantizing a BF16 model to MXFP8 or NVFP4. One would hope that the mismatch introduced by changing the engine is smaller in magnitude than the mismatch introduced by quantizing the model.

In practice, this is not necessarily true, and we found instances where engine mismatch was greater than quantization mismatch for some FP8 checkpoints. For QAD, this is particularly concerning. We are training the model to compensate for quantization error, but the error it sees during training may differ from the error it needs to correct during deployment.

To measure this more concretely, we created a quantization error alignment metric.

Let $z_{E,q}$ and $z_{E,0}$ denote the quantized and original BF16 logits in engine $E$, where $A$ is the training actor and $R$ is the rollout engine:

We used a centered, teacher probability-weighted cosine. With $p$ denoting the original BF16 teacher distribution from the training actor, we construct:

To motivate this weighting, consider a perturbation $\epsilon$ around teacher logits $z$, with $p=\operatorname{softmax}(z)$ and $q=\operatorname{softmax}(z+\epsilon)$. For small perturbations, a second-order Taylor approximation gives:

We use the same teacher distribution to weight both engines’ errors, putting them in a common local KL geometry.

We center the logit errors to remove shifts that do not change the output distribution. For example, if quantization adds 0.2 to every logit, the probabilities remain unchanged, but every entry of the logit-error vector is 0.2. Without centering, two engines could appear to have aligned errors simply because they introduce similar offsets. Subtracting each error vector’s teacher-weighted mean removes this component. We compute the resulting cosine at each scored position and average over positions.

A cosine near one means the engines introduce errors in similar directions. It does not mean the errors are small, or that the engines would produce the same parameter gradients.

Our initial FP8 quantization error alignment was 0.407. Within-engine BF16-to-FP8 KL was similar: 0.0027 for the training actor and 0.0028 for the rollout engine. But KL between the engines was 0.00241 in FP8, compared with 0.000385 in BF16: approximately 6.3 times larger. The engines introduced similar amounts of quantization error in substantially different directions.

We traced these discrepancies to four recurring sources:

1. Batching and kernel selection. GPU libraries can select different kernels based on batch shape and parallelism layout. For example, a singleton call may use split-K to expose more parallel work by dividing each dot product’s input-feature dimension into groups. Computing and combining those partial results can change accumulation order and rounding compared with the kernel selected for a larger batch. [13]⌝
2. Worker assignment and chunking. A long sum can be divided into partial sums, computed independently, and then combined. Changing which terms each worker accumulates, or how the subtotals are combined, changes the intermediate values and rounding.
3. Operand format and compute precision. BF16 and FP8 GEMMs can use different tensor-core instructions and internal arithmetic even when their operands represent the same numerical values. In the inspected kernels, BF16 instructions incorporate groups of 16 products, while FP8 instructions incorporate groups of 32. Both use FP32 accumulators.
4. Casts and fusion boundaries. Some implementations explicitly convert intermediate results to lower-precision formats before the next operation consumes them. One implementation may round a product or residual sum before consuming it, while another retains the higher precision value. This is common for fused operations in particular.

Matching these paths required changes to projections, partial-sum rounding, normalization, and attention implementations. In our final diagnostic configuration, FP8 trainer engine to rollout engine KL was 0.0000459 and error alignment was 0.924. This required giving up optimizations used by the original paths. We are therefore more interested in making QAD resilient to engine mismatch than requiring identical numerical execution.

Quantization-Aware Distillation Experimental Results

We used Qwen 3.8 27B on the same dataset, generating traces from the original BF16 teacher and training the quantized student on those traces. In this initial study, we quantized both attention and MLP projections, while preserving the LM head and linear-attention A/B gates.

The standard distillation objective is forward KL:

Here, $h$ is a prefix from a teacher-generated trace, $p_0$ is the frozen original teacher distribution, and $q_{A,\theta}$ is the quantized student distribution in the training actor.

We tested two approaches:

1. Normal QAD. Calculate the distillation loss using the training actor’s student probabilities, then backpropagate through its QDQ forward pass.
2. Deployment-informed QAD. Use the rollout engine’s student probabilities to calculate the error, then propagate that error through the training actor’s differentiable path.

The second approach is a gradient surrogate. In standard distillation, we have:

We’d like to rewrite this gradient such that we have surrogate $\widetilde{\mathcal L}$ where:

In other words, our logit update should come from the difference between rollout engine probabilities versus teacher probabilities rather than training actor probabilities versus teacher probabilities. The scalar loss that achieves this gradient is:

Here, $z_A$ is the training actor’s student logit vector before softmax, with $q_A=\operatorname{softmax}(z_A)$. The operator $\operatorname{stopgrad}$ treats the entire correction $q_R-q_A$ as constant during backpropagation. We average this loss over scored positions. Its scalar value is a surrogate, not the deployed model’s KL.

This supplies the backward pass with the deployed model’s error. It still uses the training actor’s parameter sensitivity, so it is not the exact gradient of the deployed model.

For each approach, we ran both top-500 forward KL and one-hot cross entropy, both using teacher-forced sequences. We trained on approximately 40 million tokens.

Normal top-500 QAD produced the largest reduction in deployed KL: held-out KL against the original BF16 model fell from 0.03485 to 0.02064, a reduction of approximately 41%. The reductions were similar in and out of sample. Deployment-informed training produced a smaller KL reduction, and one-hot training reduced KL less than top-500 distillation.

In a subsequent experiment, we combined MLP-only quantization with a longer training schedule, running normal top-500 QAD for six epochs over 896 teacher-generated traces. The deployed model reached an original-teacher top-500 KL of 0.00761 in sample and 0.00916 on held-out traces.

These results showed that QAD could improve the deployed model despite engine mismatch. Normal QAD achieved the lowest KL, while deployment-informed training performed better in some task-reward comparisons below.

One hypothesis we had was that QAD might reduce quantization error across engines by pushing weights towards representations that are more resilient to rounding generally. You can measure this via reconstruction error: you quantize and then de-quantize the weights, computing the introduced residual from the original weights from this operation. Interestingly, following QAD, reconstruction error increased by approximately 0.011%. Thus, it seems that QAD was able to reduce KL by adapting model behavior to account for quantization error, rather than reducing weight matrix reconstruction error directly. Note that weight matrix reconstruction error does not measure activation quantization error.

Lower KL Often Does Not Recover Accuracy

Across our trained checkpoints, reductions in held-out KL computed on a set of 180 held-out teacher traces did not consistently correspond to higher task reward on 800 held-out evaluation examples.

Though various recipes demonstrated significantly recovered task accuracy, the checkpoint with the lowest KL was often not the checkpoint with the highest reward. In one instance, normal top-500 QAD reduced deployed KL by approximately 41%, but its mean reward was 0.348, almost identical to NVFP4 without QAD.

Deployment-informed one-hot training produced the highest measured reward among the trained checkpoints, but it only reduced KL by approximately 10%. It scored 0.390, compared with 0.348 for the untrained NVFP4 and 0.371 for the original BF16 base. It’s surprising that such a small reduction in KL can recover so much lost accuracy, and that checkpoints with markedly more improved KL did not experience comparable or better accuracy recovery.

Why might the rankings differ? KL measures how much the output distribution changes, but it does not account for which changes matter to the task. Moving probability between two equally unsuccessful outcomes can change KL without changing reward. Moving probability toward a successful outcome can improve reward while increasing divergence from the original model.

We view reducing KL as a useful objective for preserving model behavior, but its relationship with task performance can be noisy. Future work along this direction will involve studying other metrics for quantization that could proxy for task performance.

What’s Next?

Keeping sensitive components in higher precision and performing QAD provides a practical approach for recovering accuracy lost to quantization. These techniques significantly reduced the deployed KL of quantized models, despite training-engine mismatch, and some recipes recovered task reward to the range of the original BF16 model. However, the relationship between KL and recovered reward was inconsistent.

We’re excited to continue working on low-precision training techniques. We’re investing in accuracy-preserving methods like QAD, as well as techniques meant to better prepare the model for quantization during capabilities training, like QAT. Additionally, we’re exploring adopting low-precision computation more broadly in the training process, such as in backward passes and optimizer state.

If these problems are exciting to you, we’re hiring! Please reach out at hiring@appliedcompute.com⌝.

Appendix: Quantization after SFT and RL

We also checked whether behavior learned during SFT or RL was especially sensitive to quantization. We trained Qwen 3.8 27B on the same dataset using GRPO, then generated roughly 800 training rollouts from that checkpoint for two ordinary SFT runs, at learning rates of $5\times10^{-6}$ and $5\times10^{-5}$. The SFT runs were not quantization-aware. We applied MXFP8 and NVFP4 after training and evaluated on the same 800 held-out examples.

We tried two learning rates because small weight updates may be more easily lost to rounding during quantization. Motivated by prior work on training dynamics and quantization robustness⌝, we wanted to check whether the lower-learning-rate SFT checkpoint was more sensitive.

The GRPO checkpoint retained most of its improvement: mean reward fell from 0.678 in BF16 to 0.644 in NVFP4, a 5.0% relative drop, compared with 6.3% for the original model. Neither SFT checkpoint showed a clear NVFP4 penalty. At the lower learning rate, reward was 0.655 in BF16 and 0.661 in NVFP4; at the higher learning rate, it was 0.627 and 0.643. Both paired 95% intervals for the NVFP4 penalty included zero.

The chart compares each quantized checkpoint with its own BF16 version on the same scored examples. Positive values mean reward decreased. Intervals use a difficulty-stratified bootstrap over evaluation examples, clustering the two GRPO rollouts per example.

We did not reproduce unusually severe post-training quantization degradation in this experiment. This is one model and one task, without a training-seed sweep; it does not establish that SFT or RL generally improves quantization robustness.

References

1. A Visual Guide to Quantization⌝. Maarten Grootendorst (2024).
2. “Give Me BF16 or Give Me Death”? Accuracy-Performance Trade-Offs in LLM Quantization⌝. Kurtic et al. (2024).
3. Ampere Tuning Guide⌝. NVIDIA.
4. Using FP8 and FP4 with Transformer Engine⌝. NVIDIA.
5. Introducing NVFP4 for Efficient and Accurate Low-Precision Inference⌝. NVIDIA (2025).
6. Block-Scaled Matrix Multiplication⌝. Triton project.
7. Evaluating Quantized Large Language Models⌝. Li et al. (2024).
8. AutoQuantize: A Fast Automatic Mixed-Precision Assignment⌝. NVIDIA Model Optimizer Team (2026).
9. Quantization-Aware Distillation for NVFP4 Inference Accuracy Recovery⌝. Xin et al. (2026).
10. Quantization-Aware Training and Distillation⌝. NVIDIA.
11. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation⌝. Bengio, Léonard, and Courville (2013).
12. Pretraining Large Language Models with NVFP4⌝. NVIDIA et al. (2025).
13. Defeating Nondeterminism in LLM Inference⌝. Thinking Machines Lab (2025).
14. Lecture Notes on Statistics and Information Theory⌝. John Duchi (2023), §§2.1.2 and 2.2.4.