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
# 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.

-