Quantization-Aware Training (QAT)

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Problem: Non-Differentiable Rounding
5. Straight-Through Estimator (STE)
6. Fake Quantization Nodes
7. Training Procedure
8. LSQ: Learned Step Size Quantization
9. QAT vs PTQ: Quality Comparison
10. QAT for LLMs
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Explain why quantization breaks gradient computation and how STE fixes it.
2. Implement fake quantization nodes in a training graph.
3. Derive LSQ's learnable scale parameter with proper gradient scaling.
4. Compare QAT vs PTQ quality at various bit-widths.
5. Evaluate the cost-benefit of QAT for LLM deployment.

Notation

• STE — straight-through estimator
• $\lfloor\cdot\rceil$ — round to nearest integer
• FQ — fake quantization operation

Core Intuition

Post-training quantization (PTQ) rounds a pre-trained model's weights — the model never "practiced" being quantized. QAT simulates quantization DURING training: the model learns weights that are robust to rounding. It's like practicing a speech with a bad microphone — you learn to articulate more clearly. QAT typically recovers 0.5-1% accuracy that PTQ loses, especially at low bit-widths (INT4, INT2).

The Problem: Non-Differentiable Rounding

Quantization involves rounding: $q = \lfloor w/s \rceil$.

Rounding derivative: $\partial\lfloor x \rceil / \partial x = 0$ almost everywhere (piecewise constant).

Consequence: If we insert quantization into the forward pass, gradients are zero → weights never update → training fails.

Need: An approximate gradient that flows through the quantization operation.

Straight-Through Estimator (STE)

Bengio et al. (2013): Pretend the rounding function is the identity during backpropagation:

where $c$ is the clipping threshold (gradients zeroed for weights outside quantization range).

Justification: The STE gradient is biased but provides a useful direction. The model learns weights that minimize loss AFTER rounding, even though the gradient doesn't perfectly account for rounding.

Why it works: Rounding errors are small ($\leq s/2$). The gradient direction is approximately correct even if the magnitude is slightly off.

Fake Quantization Nodes

Implementation: Insert quantize-then-dequantize operations in the forward graph:

Forward pass: Uses quantized values (simulates inference). Backward pass: STE gradient flows through FQ nodes.

Placement: After each weight tensor and (optionally) after each activation:

• Weight FQ: Between weight and matmul.
• Activation FQ: After each layer's output (before next layer's input).

During training: Model sees quantized forward pass but receives gradients. Learns to place weights at quantization grid points.

Training Procedure

Standard QAT pipeline:

1. Start from a pre-trained FP32 model.
2. Insert FQ nodes for all target layers.
3. Calibrate initial quantization parameters (scale, zero-point) using a few batches.
4. Fine-tune with STE for 10-20% of original training steps.
5. Use small learning rate (1/10 to 1/100 of original).
6. Final model: freeze quantization parameters, export as quantized model.

Key hyperparameters:

• Learning rate: $10^{-5}$ to $10^{-4}$ (much smaller than pre-training).
• Training steps: 10K-100K (not full pre-training length).
• Warmup: Start with FP32, gradually enable quantization.
• Batch size: Same as fine-tuning (not pre-training scale).

LSQ: Learned Step Size Quantization

Problem: Fixed quantization parameters (scale $s$) may not be optimal. Why not learn them?

LSQ (Esser et al., 2020): Make scale $s$ a trainable parameter:

Gradient for $s$:

Using STE for the round operation:

Gradient scaling: Scale the $s$-gradient by $1/\sqrt{n \cdot q_{\max}}$ to normalize across layers with different sizes.

Result: 0.5-1% better than fixed-scale QAT, especially at 4-bit and below.

QAT vs PTQ: Quality Comparison

Key insight: At INT8, PTQ (especially GPTQ/AWQ) is sufficient — QAT not worth the cost. At INT4 and below, QAT provides meaningful improvement. At INT2-3, QAT is the only viable approach.

QAT for LLMs

Challenge: Full QAT of a 70B model requires pre-training-scale compute (thousands of GPU-hours).

Practical approaches:

1. QAT during fine-tuning: Quantize after pre-training; QAT during SFT/RLHF stage.

• Cost: Same as fine-tuning (hours, not weeks).
• Quality: Almost as good as full QAT for the target task.

2. QLoRA + QAT: Quantize base model to INT4, add LoRA adapters, train with FQ nodes on adapters.

• Cost: Very low (only adapter parameters trained).
• Quality: Good for task-specific deployment.

3. BitNet-style (from scratch): Design architecture for 1-bit/ternary weights from the start. Train with QAT from step 0.

• Cost: Full pre-training.
• Quality: Can match FP16 at same compute budget (larger model width).

Common Pitfalls

Pitfall 1. Training QAT for too long. QAT typically converges in 10-20% of original training steps. Longer training can overfit to the quantization grid, losing generalization.

Pitfall 2. Using too large a learning rate. The model is already near-optimal; large LR destabilizes. Start at 1/100 of pre-training LR.

Pitfall 3. Not initializing scale from calibration. Random scale → random quantization → catastrophic first steps. Always calibrate scale from actual weight/activation statistics first.

Summary

• QAT trains with simulated quantization → weights learn to be robust to rounding.
• STE provides approximate gradients through non-differentiable rounding.
• Fake quantization nodes: quantize-dequantize in forward; identity in backward.
• LSQ learns optimal scale parameters (0.5-1% improvement).
• QAT vs PTQ: Essential at INT4 and below; overkill at INT8.
• For LLMs: QAT during fine-tuning is the practical sweet spot.

Exercises

Exercise 1. Derive the STE gradient for a weight $w=0.37$ quantized to INT4 symmetric with $s=0.1$. What gradient does the weight receive?

Exercise 2. For LSQ: derive $\partial\hat{w}/\partial s$ for the clipped case (weight exceeds quantization range).

Exercise 3. Compare the training cost (GPU-hours) of QAT vs GPTQ for quantizing LLaMA-2 7B to INT4. When is QAT cost-justified?

Exercise 4. Design a QAT schedule for a 13B model: specify learning rate, steps, batch size, which layers to quantize, and warmup strategy.

Exercise 5. Explain why STE is biased (the expected gradient doesn't equal the true gradient) but still leads to convergence (informal argument).