Skip to content

Quantization-Aware Training

Overview

QAT simulates quantization during training — quant-dequant ops in the forward pass — so the network learns weights that survive int8 rounding. When PTQ (Ch 06-03) costs you more than ~1% accuracy, QAT is the standard fix: it routinely recovers most of the drop.

  • Insert fake-quant nodes (FakeQuantize) between ops; forward uses rounded values, backward uses straight-through estimator.
  • Train as usual (keep optimizer/AMP), then convert() to real quantized kernels.
  • The classic flow: train fp32 → PTQ → measure → if bad → fine-tune with QAT (usually from the quantized-int readiness) → convert.
  • torch.ao.quantization supports both native and fx-based QAT.

💡 Straight-through estimator (STE): round() is non-differentiable → backward passes gradient as if identity. This is the trick that makes QAT work.

FX-Based QAT (modern path)

import torch, torch.nn as nn
from torch.ao.quantization.quantize_fx import prepare_qat_fx, convert_fx

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 8, 3, padding=1)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(8 * 8 * 8, 2)
    def forward(self, x):
        return self.fc(self.relu(self.conv(x)).flatten(1))

Prepare (QAT mode)

qconfig = torch.ao.quantization.get_default_qat_qconfig("fbgemm")
model = torch.ao.quantization.quantize_fx.prepare_qat_fx(
    Net().train(), qconfig, example_inputs=(torch.randn(1, 3, 8, 8),))
# FakeQuantize nodes now active in forward

Train normally

opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loader = [(torch.randn(8, 3, 8, 8), torch.randint(0, 2, (8,))) for _ in range(20)]

for xb, yb in loader:
    opt.zero_grad()
    loss = nn.functional.cross_entropy(model(xb), yb)
    loss.backward()
    opt.step()                    # STE: quantize-aware gradients flow

💡 Keep learning rate small; QAT is a fine-tune, not a fresh train (start from a converged fp32 model).

Convert & eval

model.eval()
model_q = torch.ao.quantization.quantize_fx.convert_fx(model)
out = model_q(torch.randn(1, 3, 8, 8))
print("converted output shape:", out.shape)

The STE in One Equation

forward:  y = round(x / s) * s
backward: dy/dx = 1        (ignore rounding derivative)

# equivalent of what FakeQuantize does:
class FakeQuant(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, s):
        return torch.round(x / s) * s
    @staticmethod
    def backward(ctx, g):
        return g, None      # STE

w = torch.tensor([1.3, 2.7, -0.5], requires_grad=True)
y = FakeQuant.apply(w, 0.5)
y.sum().backward()
print("grad (unrounded):", w.grad)   # [1., 1., 1.] — rounding ignored

When QAT Beats PTQ (real-world numbers)

Model type PTQ drop QAT drop
Vision (ResNet ImageNet) 0.5-2% ~0.2-0.5%
Detection/Segmentation heads 2-5% ~0.5-1%
Text LLM (int8) 0-2% ~0%
Outlier-heavy GNNs/transformers 5%+ ~1%

Decision: try PTQ first (cheap); escalate to QAT when the drop exceeds your SLA. QAT costs a training run — budget it.

QAT Best Practices

  1. Start from converged fp32 weights; QAT is fine-tuning (low LR, shorter schedule).
  2. Fuse modules before QAT prepare (fuse_modules / fx auto-fusion) — fake-quant after fusion.
  3. Freeze BN stats near the end or use small batches carefully.
  4. Match deployment scheme: symmetric int8 on x86/arm, fp8 on Hopper, etc.
  5. Watch for fake-quant poison: don't run QAT-prepared model in eager fp32 eval and call it quantized — always convert first.

Detecting Converged QAT

# after training, compare:
loss_at_start = -math.log(2)   # placeholder: your baseline
# heuristic: eval fp32-in-QAT-sim (model before convert) vs converted model
# they should be close (<0.5% acc gap) -> conversion is faithful

Key Takeaways

  • QAT = train with simulated quantization (FakeQuantize) via STE; convert afterwards.
  • Use prepare_qat_fx → train → convert_fx; start from a good fp32 model.
  • STE makes rounding trainable; low LR fine-tuning recovers accuracy.
  • Always convert before measuring int8 accuracy; keep the PTQ path for cheap first passes.
  • QAT is the accuracy safety net when PTQ fails your SLA.