Quantization-Aware Healing: A Practical Recipe for Recovering Compressed, 4-Bit LLMs
Bakbergen Ryskulov, Iker García-Ferrero, David Montero, David Jansen, Ali Hashemi, Jezabel R. Garcia, Antonio Tiene, Román Orús
cs.CL, cs.AI, cs.LG, cs.PF
2026-08-21
QAH distills MXFP4 60B from original GPT-OSS 120B: it matches or beats its bfloat16 source on 7 of 9 benchmarks, peaks about 7 times faster than QAT, and does not collapse.
Cheap serving now usually means two cuts in a row: shrink the architecture, then quantize to 4 bits. Together they hurt reasoning, math, coding, and long context. The model has to be healed before it ships.
The default heal is quantization-aware training (QAT): fake-quantizers in the forward pass, hard-label cross-entropy on the student. In Multiverse Computing's pipeline, QAT was slow and collapsed past its peak, so a safe checkpoint needed hand-tuned early stopping. NVIDIA's quantization-aware distillation (QAD) uses a frozen full-precision copy as teacher. That copy exists when compression is quantization only. After structural compression, the only bfloat16 checkpoint is a distillation-recovered approximation of the original. Distilling from it caps the student at a degraded target.
QAH changes the teacher. The 4-bit student is distilled from the original uncompressed model, not from the recovered bfloat16 checkpoint. Teacher and student no longer share an architecture. Supervision is the teacher's output distribution, the same idea used in DistilBERT-style cross-architecture distillation.
The loss is KL at temperature 1. MXFP4 fake-quantizers sit in the student forward pass with a straight-through estimator. Teacher logits are precomputed offline and truncated to the top 100; the student never sees hard labels. Offline top-k plus a fused chunked KL drops peak intermediate memory from O(BLV) to linear in sequence length, which is what makes 32k-context healing fit on the same hardware as short-context QAT.
The pipeline has two stages. A tensor-network operator compresses GPT-OSS 120B to 60B (and 20B to 9B on an internal job). The compressed model is first recovered in bfloat16 by KL distillation from the original teacher, then re-quantized to MXFP4 with a second KL pass against that same uncompressed teacher. That second pass is QAH. Embeddings, layer norms, and selected attention blocks stay frozen; unfreezing them at higher learning rates produced checkpoints worse than unhealed MXFP4.
The QAT baseline uses the same fake-quantizers and swaps KL for next-token cross-entropy.
On GPT-OSS 120B to 60B to MXFP4, the 4-bit QAH student matches or beats its own bfloat16 source on 7 of 9 benchmarks, at about one quarter of the 16-bit weight memory and half the teacher's parameter count. LiveCodeBench is 66.5 versus the 120B teacher's 66.0; the paper reads that as a match.
| Setting | Metric | Number |
| QAH 60B vs own bf16 | AA-LCR | 42.7 vs 35.3 (+7.4) |
| same | AIME 2025 | 76.3 vs 70.7 (+5.6) |
| same | Aider | 40.9 vs 38.2 (+2.7) |
| same | τ²-bench | 61.7 vs 59.4 (+2.3) |
| same | MMLU-Pro / SciCode | -0.2 / -1.4 |
| QAH 60B vs 120B teacher | LiveCodeBench | 66.5 vs 66.0 |
| same | GPQA Diamond | 67.4 vs 69.0 |
| same | AA-LCR | still -7.3 |
QAH versus QAT is measured on the 20B-to-9B MXFP4 line, averaging MMLU-Pro, LiveCodeBench, and GPQA Diamond. QAH hits 54.9 in about 100 steps and stays within about two points through 1200 steps. QAT hits a similar 54.6 only at step 700, then falls to about 36 by step 1200, a drop of nearly 19 points. That is about 7 times fewer steps to peak. At 60B, QAH peaks in about 400 steps.
The distributed backend is a real hyperparameter. Across eleven QAT configs on the 120B student, the best FSDP2 run reaches 73.74 on GPQA Diamond; no DeepSpeed ZeRO-3 run exceeds 65.15. MMLU-Pro is 70.95 versus 69.45; AIME 2025 is 80.00 versus 76.67.
This is a deployable recipe for compress-then-quantize, without a multi-week search. Quantization becomes a second chance to apply teacher supervision. The 4-bit artifact can be at least as strong as its 16-bit source and cheaper to serve. Hypernova-60B is released Apache 2.0; the numbers here are pipeline measurements, and the public weights include extra training.
If the model is quantized but not structurally compressed, standard QAD remains the closer recipe. QAH is for the compound case.
The head-to-head that would prove the teacher swap is missing: QAH versus QAD from the recovered bfloat16 checkpoint at matched settings. The recovered-teacher ceiling is an argument, not a measurement. Every score is a single run. AIME 2025 has 30 problems. The study covers one family (GPT-OSS MoE), one format (MXFP4), and one Nemotron plus SmolTalk mix. The compression operator is proprietary, so transfer to pruning, SliceGPT, or low-rank methods is untested. GPT-OSS weights are already largely MXFP4; the gain is from an uncompressed teacher, not a higher-precision one. The DeepSpeed gap is blamed on mixed-precision communication, without a pinpoint diagnosis.