Lab 02 — Quantization-Aware Training (QAT)

Phase: 03 — Quantization & Model Compression | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours

QAT trains with simulated quantization noise — the model learns to be robust to INT8/INT4 precision.

What you build

  • FakeQuantize — asymmetric per-tensor fake-quant with Straight-Through Estimator (STE)
  • QATLinear — Linear layer with fake-quantized weights
  • wrap_model_for_qat — replace all Linear layers with QATLinear (deepcopy safe)
  • train_qat — AdamW training loop on QAT model
  • convert_to_int8 — fold scale/zero_point into INT8 weights post-training

Key concepts

ConceptWhat to understand
STEx_out = x + (quant(x) - x).detach() — gradients flow through
FakeQuantizeQuantize + dequantize in forward; model learns to minimize error
INT8 conversionAfter QAT, actual INT8 weights match dequantized FakeQuant output
Accuracy vs PTQQAT typically +1–3% accuracy over PTQ at same bit width

Files

FilePurpose
lab.pyStubs with TODO markers
solution.pyComplete working implementation
test_lab.pypytest suite
requirements.txtDependencies

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v