🎓 A1Learner
🔍

PyTorch Quantization-Aware Training

Quantization-aware training trains a model with fake quantization inside its forward pass. Fake quantization rounds values to int8 steps and then expands them back to floats.

Rounding to a few levels

A coarse grid makes the rounding visible here. Five levels pull each weight to its nearest step.

import torch

def fake_quant(w, levels=5):
    scale = w.detach().abs().max() / (levels // 2)
    return torch.round(w / scale).clamp(-(levels // 2), levels // 2) * scale

w = torch.tensor([-0.30, 0.10, 0.55, 0.90])
print([round(v, 2) for v in fake_quant(w).tolist()])

Rounding alone blocks gradients

The round function has a gradient of zero almost everywhere. A loss built on rounded weights gives the weight a zero gradient.

import torch

w = torch.tensor([0.4], requires_grad=True)
y = torch.round(w * 127) / 127
loss = y.sum()
loss.backward()
print(w.grad.item())

The straight-through fix

The detach trick adds the rounding gap as a constant. The same loss now hands the weight a gradient of one.

import torch

w = torch.tensor([0.4], requires_grad=True)
rounded = torch.round(w * 127) / 127
fake = w + (rounded - w).detach()
loss = fake.sum()
loss.backward()
print(w.grad.item())

Example

import torch
import torch.nn.functional as F

def fake_quant(w):
    scale = (w.detach().abs().max() / 127.0).clamp(min=1e-8)
    rounded = torch.round(w / scale).clamp(-127, 127) * scale
    return w + (rounded - w).detach()

class FakeQuantLinear(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 1)

    def forward(self, x):
        weight = fake_quant(self.linear.weight)
        return F.linear(x, weight, self.linear.bias)

torch.manual_seed(0)
x = torch.randn(60, 4)
y = (x[:, 0] * 2 - x[:, 1] + 0.5).reshape(-1, 1)

def train(make):
    torch.manual_seed(7)
    model = make()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
    loss_fn = torch.nn.MSELoss()
    for step in range(120):
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()
    return model, loss.item()

plain, plain_loss = train(lambda: torch.nn.Linear(4, 1))
aware, aware_loss = train(FakeQuantLinear)
with torch.no_grad():
    gap = (plain(x) - aware(x)).abs().mean().item()
print(f"Plain final loss: {plain_loss:.4f}")
print(f"Fake-quant final loss: {aware_loss:.4f}")
print(f"Mean output gap: {gap:.4f}")

Key points

  • Scale: the gap between two neighbouring int8 levels for a tensor
  • Straight-through: gradients pass a rounding step as if nothing changed
  • Use w + (rounded - w).detach() to round in forward and keep gradients flowing
  • Only the weights are rounded in this small version, and the bias stays float
  • A model trained with rounding in the loop is ready for rounding later

In this example, two models fit the same seeded regression data. Both models start from the same seed, so their first weights match. The plain model is one Linear layer trained for 120 steps. The aware model wraps the same layer and rounds its weight on every forward pass. Its rounding uses 127 levels and the straight-through trick from the snippets. The closing lines print both final losses and the mean gap between their outputs. The fake-quant loss lands at 0.0003, and the mean output gap stays at 0.0179.

Note

Library QAT adds observers, which are helpers that watch value ranges in training

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does fake quantization use the detach trick?