Source-linked AI summary

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

arXiv:2608.20953v1cs.CLcs.AIcs.LGcs.PF

TL;DR

Structurally compressed and 4-bit quantized LLMs lose capabilities needed for deployment, while QAT can be slow and unstable. The paper introduces QAH, which distills the quantized student from the original uncompressed model, and reports strong benchmark quality with faster, more stable training than QAT.

  • Problem

    Structural compression followed by 4-bit quantization degrades LLM reasoning, mathematics, coding, and long-context behavior, while the standard QAT recipe is slow and operationally fragile.

  • Method

    QAH heals the compressed, quantized student by distilling its output distribution from the original uncompressed model rather than a recovered bfloat16 checkpoint.

  • Results

    On GPT-OSS 120B→60B→MXFP4, QAH matches or beats its bfloat16 source on 7 of 9 benchmarks and converges about 7× faster than matched QAT.

  • Takeaways & Limitations

    QAH provides a deployable healing recipe that remains stable under continued training and does not require hand-tuned early stopping.

  • Takeaways & Limitations

    The paper does not directly measure QAH against QAD distilled from the recovered bfloat16 checkpoint, so the recovered-teacher ceiling remains motivated rather than head-to-head tested.

Abstract

from arXiv · show

Serving large language models cheaply increasingly means shipping models that are both structurally compressed to a fraction of their parameters and quantized to 4 bits. Together these steps degrade reasoning, mathematics, coding, and long-context behavior enough to require a recovery, or healing, stage before deployment. The default recipe, quantization-aware training (QAT), re-fits the compressed, quantized model to hard labels; in our pipeline it converged slowly and collapsed past its peak. We adopted Quantization-Aware Healing (QAH) instead. Because a structurally compressed model is never independently trained at full precision, its bfloat16 checkpoint is a distillation-recovered approximation of the original; QAH distills the 4-bit student directly from the original, uncompressed model. On a GPT-OSS 120B to 60B to MXFP4 pipeline, the QAH student matches or beats its bfloat16 source on 7 of 9 benchmarks at roughly 4 times less weight memory and half the teacher's parameter count, and is released open-weight as Hypernova-60B. Against a matched QAT baseline it reaches a comparable peak about 7 times faster and stays stable under continued training, without hand-tuned early stopping. We also report deployment lessons, including a large, reproducible quality gap between distributed-training backends. Our aim is a recipe deployable without a multi-week hyper-parameter search.

1 Introduction

Structural compression followed by 4-bit quantization reduces LLM capabilities, motivating a healing stage. QAH distills the compressed, quantized student from the original uncompressed model and improves quality, speed, and deployment stability over QAT.

  • Motivation: Structural compression and 4-bit quantization reduce reasoning, mathematical problem-solving, code-generation, and long-context capabilities, requiring recovery before deployment.The paper calls this recovery step healing.
  • Limitations of QAT: QAT converges slowly and can collapse after its peak, making safe deployment dependent on early stopping against a held-out signal.QAT trains with fake quantizers and a cross-entropy task loss.
  • QAH: QAH distills the quantized student from the original uncompressed model rather than the recovered bfloat16 checkpoint or hard labels.The original model provides a stronger teacher, while teacher supervision uses an architecture-agnostic output distribution.
  • Results: On GPT-OSS 120B→60B→MXFP4, QAH matches or exceeds its bfloat16 source on 7 of 9 benchmarks and reaches the original 120B model on LiveCodeBench.The student uses roughly 4× less weight memory and half the teacher’s parameter count.
  • Deployment: QAH converges about 7× faster than matched QAT, remains stable under continued training, and avoids hand-tuned early stopping.The paper also reports a large, reproducible quality gap between distributed-training backends.

2 Background and Related Work

The background distinguishes structural compression, quantization, QAT, and knowledge distillation. QAH combines quantization-aware training with cross-architecture distillation from the original pre-compression model for the compress-then-quantize setting.

  • Low-precision formats: Microscaling formats use block-wise scaling factors to reduce memory and compute while preserving most of the bfloat16 forward pass.MXFP4 quantizes weights only and underlies the public GPT-OSS release.
  • Quantization-aware training: QAT adapts weights to quantized representations with fake quantizers but forces students to reacquire behavior already present in the original model.The paper uses QAT as its baseline.
  • Knowledge distillation: Knowledge distillation trains students to match teacher output distributions and does not require teacher and student architectures to match.QAH relies on this cross-architecture property because its teacher is uncompressed and its student is compressed and quantized.
  • Structural compression: Structural compression modifies architecture through methods such as pruning, sparsity, dimension slicing, or low-rank decomposition.These changes make the recovered checkpoint an unsuitable QAD teacher according to the paper’s motivation.
  • Quantization-aware distillation: Standard QAD uses the original bfloat16 model as a frozen teacher, but presupposes architectural identity when compression is quantization-only.QAH targets cases where structural compression precedes quantization and teacher and student architectures differ.

3 Method

The paper defines compound compression as structural reduction followed by low-precision quantization, then introduces QAH, which distills the quantized student from the original uncompressed teacher. This preserves access to the teacher’s behavior while accommodating architectural differences and long-context training constraints.

  • 3.1 The Compound-Compression Regime: Compound compression applies structural reduction and then MXFP4 quantization to the original model before deployment.The pipeline is MFP4 = Q(R(S(MHP))), where S compresses architecture, R recovers the compressed model in bfloat16, and Q quantizes it.
  • 3.1 The Compound-Compression Regime: The compressed intermediate is never independently trained at full precision, so its recovered bfloat16 checkpoint remains tethered to the original teacher and its capacity ceiling.Recovery uses KL distillation from MHP rather than training the compressed architecture from scratch.
  • 3.2 Quantization-Aware Healing: QAH replaces the recovered bfloat16 checkpoint with the original uncompressed model as teacher for the MXFP4 student.Teacher and student may have different architectures; the student is supervised through the teacher’s output distribution.
  • 3.2 Quantization-Aware Healing: QAH minimizes teacher–student KL divergence using straight-through MXFP4 fake-quantizers, masked next-token positions, and a healing corpus.Teacher logits are precomputed offline and truncated to the top 100 logits; the student sees no hard labels.
  • 3.3 Quantization-Aware Training Baseline: QAT uses the same fake-quantizers but optimizes cross-entropy on input–label pairs instead of conditioning on the original teacher during training.Unlike QAH, QAT must reacquire behavior from hard labels.
  • 3.2 Quantization-Aware Healing: At 32k context, offline top-k teacher logits and fused chunked KL reduce the prohibitive O(BLV) peak memory of direct implementation.The loss and gradient are accumulated block-by-block along the sequence.

4 Experimental Setup

The experiments compress GPT-OSS models, recover them with teacher distillation, and then re-quantize them under QAH or QAT. Training uses mixed-domain data, distributed H200 hardware, and evaluations spanning knowledge, reasoning, coding, tool use, and long-context tasks.

  • Models and pipeline: The pipeline compresses GPT-OSS 120B to 60B and 20B to 9B, then recovers and re-quantizes each model to MXFP4.The quantization stage distills from the uncompressed teacher for QAH or uses cross-entropy for QAT.
  • Data: Training data combines NVIDIA Nemotron and SmolTalk 2 across general-domain, science, coding, mathematics, safety, and reasoning tasks.
  • Training: All runs use eight NVIDIA H200 nodes with FSDP2 parameter sharding; QAH uses 32k sequences, batch size 64, learning rate 5×10^-6, and 400 steps.Quantization-sensitive submodules are frozen to limit destructive drift.
  • Evaluation: Evaluation covers nine benchmarks spanning general knowledge, science, mathematics, instruction following, scientific and agentic coding, tool use, and long-context reasoning.The suite includes MMLU-Pro, GPQA Diamond, AIME 2025, IFBench, SciCode, LiveCodeBench, τ 2-bench, Aider, and AA-LCR.

5 Results

QAH produces a compact 4-bit student that generally matches or improves on its bfloat16 source while retaining much of the teacher’s capability. Compared with QAT, it reaches a similar peak faster and remains stable during extended training.

  • 5.1 Quantization as a Second Training Opportunity: 7 of 9 benchmarks show the MXFP4 QAH student matching or beating its bfloat16 source.The exceptions are MMLU-Pro, down 0.2 points, and SciCode, down 1.4 points.
  • 5.1 Quantization as a Second Training Opportunity: +7.4 on AA-LCR, +5.6 on AIME 2025, +2.7 on Aider, and +2.3 on τ 2-bench compare the QAH student with its bfloat16 source.These gains are reported as benchmark score differences between the MXFP4 and bfloat16 60B checkpoints.
  • 5.1 Quantization as a Second Training Opportunity: The quantization stage adds a second KL-distillation pass from the original 120B distribution, rather than demonstrating that 4-bit representations are inherently superior.This extra supervision is why the deployable 4-bit artifact can be at least as capable as its bfloat16 source.
  • 5.1 Quantization as a Second Training Opportunity: At half the teacher’s parameter count and 4-bit precision, QAH reaches the 120B model on LiveCodeBench and comes within 1.6 points on GPQA Diamond.The largest residual gap is −7.3 on AA-LCR, the extreme long-context benchmark.
  • Efficiency: Roughly 4× less weight memory and half the teacher’s parameter count provide the principal serving-efficiency gains.The reported reductions concern the QAH model relative to the bfloat16 student and the 120B teacher, respectively.
  • 5.2 QAH vs. QAT: Cost and Stability: 54.9 for QAH versus 54.6 for QAT is a comparable peak, but QAH reaches it in about 100 steps instead of roughly 700.The comparison uses average performance on MMLU-Pro, LiveCodeBench, and GPQA Diamond for GPT-OSS 9B quantized to MXFP4.
  • 5.2 QAH vs. QAT: Cost and Stability: QAH stays within about two points of its peak through 1200 steps, whereas QAT loses nearly 19 points after its peak.The authors attribute the asymmetry to KL distillation anchoring the student versus cross-entropy continuing to push toward hard labels.

6 Deployment Lessons

Deployment depends not only on the healing recipe but also on backend choice, quantization-sensitive freezing, and memory-efficient long-context training. These operational choices determine optimization quality, stability, and whether 32k-context healing fits the hardware envelope.

  • Backend choice: The distributed backend is a tuned hyperparameter: FSDP2 and DeepSpeed ZeRO-3 produced a large, consistent optimization-quality gap.The sweep covered context length, batch size, node count, and learning rate; the authors default to FSDP2.
  • Quantization-sensitive modules: Freezing embeddings, layer norms, and selected attention components was necessary for stable healing under both QAT and QAH.Unfreezing layer norms and embedding projections, especially at higher learning rates, produced checkpoints worse than the unhealed MXFP4 model.
  • Long-context healing: 32k-context QAH fit the same hardware as short-context QAT by using chunked-KL loss, which reduces peak intermediate memory from O(BLV) to O(BcV).The implementation remains bit-identical to the dense loss, making long-context healing tractable within a fixed GPU-memory envelope.

7 Limitations and Open Questions

The evaluation is constrained by an unrun direct QAD comparison, single-run and small-sample measurements, limited model and format coverage, and a proprietary compression operator. These boundaries leave transferability and the central teacher comparison incompletely measured.

  • Open comparisons: The direct QAD baseline was not run, so QAH’s recovered-teacher-ceiling argument remains well-motivated but not directly measured.The missing experiment is QAH versus QAD from the recovered bfloat16 checkpoint under matched configurations.
  • Evaluation uncertainty: Every number is from a single run without seed variance or confidence intervals, and some benchmarks are small, so individual deltas are indicative rather than significant.Comparisons within one or two points, especially LiveCodeBench versus the teacher, are read as matches rather than wins.
  • Scope: All experiments use GPT-OSS MoE transformers, MXFP4 quantization, and one Nemotron + SmolTalk mixture; other model families and formats were not evaluated.The untested formats include NVFP4, INT4, and FP8, while untested families include Llama, Qwen, and Mistral.
  • Scope: The structural-compression operator is proprietary, and transfer to layer pruning, SliceGPT, or low-rank decomposition was not verified.The authors state that QAH is compression-method agnostic in principle but have not tested these alternative approaches.

8 Conclusion

The paper presents QAH as a practical recovery recipe for structurally compressed, 4-bit LLMs. It reports lower serving cost, faster and more stable healing than QAT, and deployment lessons alongside an explicit unrun comparison.

  • Conclusion: QAH distills a structurally compressed, 4-bit student from the original uncompressed model rather than a recovered full-precision checkpoint.The recipe combines teacher supervision with quantization-aware training for the compressed student.
  • Conclusion: On GPT-OSS 120B→60B→MXFP4, the 4-bit QAH model matches or beats its bfloat16 source on 7 of 9 benchmarks at roughly 4× lower weight memory and half the parameter count.It reaches the 120B model on LiveCodeBench.
  • Conclusion: Against a matched QAT baseline, QAH converges about 7× faster and does not collapse under continued training, avoiding hand-tuned early stopping.The conclusion also reports a large, reproducible quality gap between distributed-training backends.
  • Conclusion: The paper explicitly identifies the missing direct QAD comparison as an important qualification on its recovered-teacher claim.That comparison was not run at matched configuration.

A QAT Training-Systems Sweep

The QAT systems sweep shows that distributed-training configuration materially affects optimization quality on the 120B student. FSDP2 was the only configuration to reach the released MXFP4 baseline on GPQA Diamond.

  • QAT Training-Systems Sweep: Only FSDP2 configuration R11 reached the released 120B MXFP4 baseline on GPQA Diamond; every DeepSpeed configuration remained 3.5–14.2 points below it.The eleven-run sweep varied backend, effective sequence length, learning rate, batch size, cluster size, and layer-norm freezing strategy.
Loading 2608.20953v1…