Quantization-Aware Training: How QAT Recovers Lost Accuracy in 4-Bit LLMs

Every LLM developer faces the same tradeoff: run models at full precision for accuracy, or quantize them down to 4-bit for speed and cost — and accept the quality hit. Until now.

Quantization-Aware Training (QAT) is changing that equation. By simulating quantization during training, QAT teaches models to adapt to low-precision arithmetic before deployment, recovering accuracy that post-training quantization (PTQ) would otherwise destroy. The results are striking: PyTorch reports 96% accuracy recovery on hellaswag for Llama 3, Unsloth demonstrates 1-3% net performance gains on benchmarks like GPQA and MMLU Pro, and Google's Gemma 4 QAT models now run on as little as 3GB of RAM with near-original performance.

This deep dive covers what QAT is, how it works under the hood, the frameworks making it accessible, and why it matters for the future of on-device AI.

The Problem: Why Naive Quantization Destroys Accuracy

Post-training quantization (PTQ) is the simplest way to compress a model. Take the weights, find the maximum absolute value, compute a scaling factor, and round everything to integers. For 8-bit quantization, the math is straightforward:

PTQ dramatically reduces storage and inference costs. But at 4-bit precision, the rounding errors accumulate fast. A model fine-tuned in BF16 has weights optimized for that precision — quantizing them afterward is like taking a photograph printed at 300 DPI and shrinking it to 50 DPI. Details get lost, and the model's ability to produce accurate outputs degrades proportionally.

The degradation is especially severe at 2-bit and 3-bit precision. PyTorch's experiments with Llama 3 at 2-bit weight-only quantization showed wikitext perplexity exploding from ~10 to over 600,000 — the model essentially stopped functioning.

The Solution: Quantization-Aware Training

QAT addresses this by making quantization part of the training process itself. Instead of applying quantization after training, QAT inserts "fake quantization" operations into the model's forward pass during training. The weights stay in high precision (BF16 or FP16), but the model computes with quantized values — rounding and clipping just like it would at inference time.

Here's the key insight: the model learns to compensate for quantization noise while it's still trainable. The gradients flow back through a straight-through estimator (STE), a technique that approximates gradients through the non-differentiable rounding operation. The model adjusts its weights knowing that those weights will eventually be quantized, so it finds solutions that are robust to quantization error.

After QAT training, a conversion step replaces the fake quantize operations with actual quantize/dequantize operators, producing an inference-ready model in the target precision. The converted model is a drop-in replacement for a PTQ model — same quantization format, same memory footprint, same kernel requirements.

How QAT Works: The Two-Step Flow

The QAT pipeline has two distinct phases, both now accessible through standard PyTorch APIs:

Step 1 — Prepare (fake quantization): Insert fake quantize operations into the model's linear layers. During the forward pass, weights are quantized to the target precision and immediately dequantized back to high precision. The model computes with approximate values, and the loss function sees the quantization error. During the backward pass, gradients flow through the STE to update the original high-precision weights.

Step 2 — Convert (real quantization): After training, replace the fake quantize operations with actual quantize/dequantize operators. The model is now truly quantized and ready for inference.

Both steps are handled by torchao, PyTorch's official quantization library. The API is simple: call prepare() before training and convert() after.

PyTorch's QAT Results: Near-Perfect Recovery

PyTorch's end-to-end QAT experiments, published in July 2024 by the TorchAO team [Link], demonstrate the technique's potential. Fine-tuning Llama 2 and Llama 3 on the C4 dataset with int8 dynamic activations and int4 grouped weights (the "8da4w" scheme), they measured accuracy recovery across multiple benchmarks.

The 2-bit results are particularly remarkable. By skipping quantization for the first 3 and last 2 layers (identified as most sensitive through prior analysis), and applying QAT to the remaining layers, PyTorch brought wikitext perplexity from 6,766 (PTQ with layer skipping) down to just 30 — a reduction of over 220×. Without layer skipping, PTQ at 2-bit saw perplexity explode to over 600,000. Without QAT, 2-bit Llama 3 was essentially unusable.

For 3-bit quantization, QAT was effective even without skipping layers, recovering 63% of hellaswag accuracy degradation and 72% of wikitext perplexity degradation compared to PTQ.

Unsloth + TorchAO: QAT Meets LoRA

Unsloth, the popular library for fast and memory-efficient fine-tuning, integrated QAT directly into its pipeline in collaboration with the PyTorch/TorchAO team [Link]. This integration unlocks a powerful combination: QAT for accuracy recovery plus LoRA for memory-efficient fine-tuning.

Unsloth supports multiple QAT schemes through the qat_scheme parameter:

• fp8-int4: FP8 activations with int4 weights
• fp8-fp8: Full FP8 quantization
• int8-int4: Int8 activations with int4 weights (8da4w)
• int4: Weight-only int4 quantization

After fine-tuning, call model.save_pretrained_torchao() to export in TorchAO's PTQ format, uploadable directly to Hugging Face. The merged QAT model runs in vLLM, Unsloth, and other inference systems.

Unsloth's QAT results on Gemma 3 models:

Crucially, QAT adds zero inference overhead. The quantized model uses the same disk space, same memory, and same compute as a PTQ model — all the benefits come during training.

NVIDIA's QAT + QAD: Adding Knowledge Distillation

NVIDIA's approach extends QAT with Quantization-Aware Distillation (QAD), combining quantization simulation with knowledge distillation [Link]. In QAD, the "student" model is the QAT model running with fake quantization, and the "teacher" is the original full-precision model. The distillation loss measures how far the quantized predictions deviate from the teacher's outputs, providing additional gradient signal for accuracy recovery.

NVIDIA demonstrated QAD on their Llama Nemotron Super model using TensorRT Model Optimizer:

NVIDIA recommends running QAT for a duration equivalent to about 10% of initial training epochs, though in practice, even less than 1% of original pre-training time is often sufficient for LLMs. Their implementation supports FP8, NVFP4, MXFP4, INT8, and INT4 formats through familiar PyTorch and Hugging Face APIs.

One key finding: NVFP4 consistently outperforms MXFP4 on visual question answering benchmarks. NVFP4's finer-grained scaling (smaller blocks and higher-precision scale factors) better preserves small signals and occasional outliers — critical for tasks like InfoVQA, where models need to detect tiny numbers and thin lines in complex charts.

Gemma 4 QAT: Running 26B Models on 15GB of RAM

Google DeepMind's Gemma 4 family includes dedicated QAT variants designed for consumer hardware [Link]. These models are trained with quantization in mind from the start, achieving ~72% lower memory usage with near-original performance.

The QAT lineup spans from ultra-compact to large-scale:

The E2B model runs on just 3GB of RAM — making it feasible to run a multimodal, hybrid-thinking LLM on a basic laptop. The 26B-A4B MoE model, with 140+ language support and 256K context, runs on 15GB — within reach of a consumer GPU like the RTX 4070 Ti Super.

The Unsloth Dynamic Method: Fixing GGUF Conversion

One practical challenge with QAT models: converting from the BF16 QAT checkpoint to llama.cpp's GGUF format is not lossless. llama.cpp uses F16 scales while QAT BF16 uses BF16 scales, and the scales are not determined optimally in the standard conversion pipeline.

Unsloth's analysis showed that naive Q4_0 conversion achieved only 24.77% byte-exactness to the original BF16 QAT lattice. Their "Unsloth Dynamic" method improved this to 99.96% byte-exactness, while also making the quantized files smaller.

For the Gemma 4 26B-A4B model, naive Q4_0 conversion scored 70.2% top-1 accuracy. Unsloth's dynamic method pushed it to 85.6% — a 15.6 percentage point improvement — while being 200MB smaller. On the E2B model, the mean Kullback-Leibler divergence (KLD) dropped from 0.05109 (naive) to 0.00173 (Unsloth dynamic), a 29× improvement.

QAT Overhead: The Tradeoffs

QAT is not free. Inserting fake quantize operations throughout the model adds computational and memory overhead during training:

• Training speed: ~34% slower than regular full fine-tuning (PyTorch's measurement on Llama 3-8B with 8da4w)
• Memory: ~2.35 GB additional per GPU with activation checkpointing (on 6× A100 80GB)
• Why: Fake quantization requires cloning weights before applying the operation, since they can't be mutated in-place

PyTorch notes that torch.compile may reduce these overheads in future versions. Unsloth's approach of combining QAT with LoRA partially mitigates the memory cost by only updating low-rank adapter weights rather than the full model.

Empirically, both PyTorch and NVIDIA found that disabling fake quantization for the first N training steps (PyTorch used 1,000 steps) led to better results — allowing the weights to stabilize before introducing quantization noise.

Practical Guide: When to Use QAT

Not every model needs QAT. The decision pipeline is straightforward:

1. Try PTQ first. Many models retain over 99.5% of original accuracy with just post-training quantization, especially at 8-bit.
2. If PTQ accuracy is insufficient at 4-bit, apply QAT. Fine-tune with fake quantization for a small fraction of your training budget.
3. If QAT still isn't enough, try QAD. Add knowledge distillation from the full-precision teacher for additional recovery.
4. For extreme quantization (2-3 bit), skip sensitive layers. The first and last transformer layers are typically the most quantization-sensitive.

The sweet spot for QAT is 4-bit quantization, where PTQ degradation is significant but QAT recovery is reliable. At 8-bit, PTQ is usually good enough. At 2-bit, even QAD struggles without layer skipping.

Deployment Paths

QAT models deploy through the same pathways as PTQ models — the key advantage is that you get better accuracy for the same format. Current deployment options include:

• vLLM: Run merged QAT models directly for high-throughput serving
• llama.cpp: GGUF format for CPU inference, with Unsloth's dynamic quants for best accuracy
• ExecuTorch: PyTorch's on-device inference runtime for iOS and Android (via XNNPACK backend)
• TensorRT-LLM: NVIDIA's optimized inference engine with NVFP4/MXFP4 support
• Unsloth Studio: Open-source web UI for local inference and training on macOS, Windows, and Linux

Looking Ahead

QAT is rapidly maturing from a research technique to a production-ready workflow. Several directions promise even better results:

• torch.compile integration: PyTorch is working on making fake quantization compatible with torch.compile, which could eliminate the 34% speed penalty
• Mixed-precision QAT: Instead of skipping sensitive layers entirely, quantize them at higher precision (e.g., 8-bit for embeddings, 4-bit for the rest)
• MX4 and newer dtypes: NVIDIA's Blackwell GPUs will drop int4 tensor core support, making formats like MXFP4 and MX4 critical for future hardware
• QAT for embeddings and KV cache: Current QAT focuses on linear layers, but embedding quantization and KV cache compression are the next bottlenecks for long-context models
• Outlier suppression: Techniques like SmoothQuant and SpinQuant from the PTQ world could be adapted for QAT to handle activation outliers more gracefully

Summary

Quantization-Aware Training is the missing link between aggressive model compression and production-quality accuracy. By simulating quantization during training, QAT allows models to learn weight configurations that are inherently quantization-friendly — recovering 60-96% of the accuracy that PTQ would otherwise destroy.

The ecosystem is now mature enough for production use. PyTorch's torchao provides the core infrastructure, Unsloth makes it accessible through LoRA-composable fine-tuning, NVIDIA extends it with knowledge distillation, and Google ships QAT-native models like Gemma 4 that run on consumer hardware.

For anyone deploying LLMs at scale — whether on cloud GPUs, edge devices, or mobile phones — QAT is no longer optional. It's the difference between a quantized model that works and one that barely functions.

Fact Check Report

Sources

• Unsloth: "Quantization-Aware Training (QAT)" — [Link]
• PyTorch Blog: "Quantization-Aware Training for Large Language Models with PyTorch" — [Link]
• NVIDIA Developer Blog: "How Quantization Aware Training Enables Low-Precision Accuracy Recovery" — [Link]
• Unsloth: "Gemma 4 QAT" — [Link]