Quantize Callback

Quantization-aware training inside a fastai fit, with the FX and the pt2e flow

Overview

QuantizeCallback runs quantization-aware training inside a fastai fit: fake-quantize modules go in before it, the converted model comes out on learn.model. The backend picks the flow — x86, qnnpack, fbgemm and onednn give back an FX quantized model, pt2e a symmetric INT8 graph.

Both arms below continue the same fine-tuned ResNet-18.

1. A shared starting point

Oxford-IIIT Pet, two classes, 64px; three epochs of fine-tuning.

print("torch", torch.__version__)

path = untar_data(URLs.PETS)
files = get_image_files(path/"images")

def label_func(f): return f[0].isupper()

dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(64), bs=64)

torch.manual_seed(0)
model = resnet18(weights=ResNet18_Weights.DEFAULT)
model.fc = nn.Linear(512, 2)

learn = Learner(dls, model, metrics=accuracy)
with learn.no_bar(): learn.fit_one_cycle(3, 1e-3)
base = deepcopy(learn.model)

torch 2.9.1+cu128
[0, 0.3727749288082123, 0.2615922689437866, 0.8917456269264221, '00:04']
[1, 0.24850593507289886, 0.21142429113388062, 0.9147496819496155, '00:04']
[2, 0.1141485646367073, 0.18102337419986725, 0.9309878349304199, '00:04']

Every accuracy on this page goes through report, which prints the count k, its n, and the Wilson 95% interval for that n.

def evaluate(model, dl):
    "Correct predictions and n, on the CPU where a quantized model runs"
    model.eval().cpu()
    correct = total = 0
    with torch.no_grad():
        for xb, yb in dl:
            correct += int((model(xb.cpu()).argmax(-1) == yb.cpu()).sum())
            total += len(yb)
    return correct, total

def report(name, k, n, z=1.96):
    "A count with its n and its Wilson 95% interval"
    p, d, centre = k/n, 1 + z**2/n, k/n + z**2/(2*n)
    half = z*((p*(1-p) + z**2/(4*n))/n)**0.5
    print(f"{name}: {k}/{n} = {p:.4f}  Wilson 95% [{(centre-half)/d:.4f}, {(centre+half)/d:.4f}]")

2. Two more epochs, with and without the callback

Seed-matched arms: both continue from the same base weights, torch.manual_seed(0) is set again right before each one, and the only difference between them is the callback.

torch.manual_seed(0)
fp32 = Learner(dls, deepcopy(base), metrics=accuracy)
with fp32.no_bar(): fp32.fit_one_cycle(2, 1e-4)

report("FP32", *evaluate(fp32.model, dls.valid))

[0, 0.07825697213411331, 0.19201131165027618, 0.9357239603996277, '00:05']
[1, 0.0505782850086689, 0.20074552297592163, 0.9350473880767822, '00:04']
FP32: 1382/1478 = 0.9350  Wilson 95% [0.9213, 0.9465]

torch.manual_seed(0)
qat = Learner(dls, deepcopy(base), metrics=accuracy)
with qat.no_bar(): qat.fit_one_cycle(2, 1e-4, cbs=QuantizeCallback(backend='x86', verbose=True))

print("model after the fit:", type(qat.model).__name__)
print(quant_spec(qat.model).as_dict())
report("x86 QAT INT8", *evaluate(qat.model, dls.valid))

Model prepared for QAT successfully
[0, 0.06549660861492157, 0.18793834745883942, 0.9357239603996277, '00:06']
[1, 0.051761697977781296, 0.19833843410015106, 0.9364005327224731, '00:06']
Converting QAT model to fully quantized model
model after the fit: GraphModule
{'backend': 'x86', 'method': 'qat', 'weight_bits': 8, 'act_bits': 8, 'qscheme': 'per_channel', 'symmetric': False, 'group_size': None, 'layer_bits': None, 'qdq_placement': None, 'skip_activations': None, 'engine': 'x86'}
x86 QAT INT8: 1385/1478 = 0.9371  Wilson 95% [0.9235, 0.9484]

The two report lines above come from seed-matched arms, and their Wilson intervals overlap, so this single run does not order them. What the cell does show is the callback running end to end: the fit hands back an FX GraphModule, scored on the same validation set as the floating-point arm, carrying the precision record quant_spec reads back — 8-bit weights and activations, per-channel, and asymmetric, which is what the FX backends produce. The pt2e arm below prints the symmetric one.

The same two models, on disk:

def size_mb(model):
    torch.save(model.state_dict(), TMP/"state.p")
    return (TMP/"state.p").stat().st_size / 1e6

print(f"FP32 state_dict: {size_mb(fp32.model):.2f} MB")
print(f"INT8 state_dict: {size_mb(qat.model):.2f} MB")

FP32 state_dict: 44.78 MB
INT8 state_dict: 11.30 MB

3. The pt2e flow

pt2e captures the model with torch.export, and a captured graph runs one batch size. fastai’s validation loader keeps a shorter last batch, so the callback refuses before training instead of during it:

refused = Learner(dls, deepcopy(base), metrics=accuracy)
try:
    with refused.no_bar(): refused.fit_one_cycle(1, 1e-4, cbs=QuantizeCallback(backend='pt2e'))
except ValueError as e:
    print(e)

Exception occured in `QuantizeCallback` when calling event `before_fit`:
    pt2e QAT captures the model at one batch size, but these dataloaders yield [6, 64]: a graph captured by torch.export only runs the batch size it was captured with. Use a batch size that divides both dataset sizes, or drop the last batch of each loader.

Drop that batch, and the fit runs:

dls.valid.drop_last = True  # 23 batches of 64; the 6-image remainder is dropped

torch.manual_seed(0)
pt2e = Learner(dls, deepcopy(base), metrics=accuracy)
with pt2e.no_bar(): pt2e.fit_one_cycle(2, 1e-4, cbs=QuantizeCallback(backend='pt2e', verbose=True))

print(quant_spec(pt2e.model).as_dict())

pt2e: capturing the model with torch.export
pt2e: prepared for QAT (W8A8, per_channel, qdq_placement=per_op)
[0, 0.0709400549530983, 0.1948516070842743, 0.932744562625885, '00:07']
[1, 0.04437658190727234, 0.19175469875335693, 0.9375, '00:07']
pt2e: converting the trained graph to a quantized one
{'backend': 'pt2e', 'method': 'qat', 'weight_bits': 8, 'act_bits': 8, 'qscheme': 'per_channel', 'symmetric': True, 'group_size': None, 'layer_bits': None, 'qdq_placement': 'per_op', 'skip_activations': None, 'engine': None}

The graph it leaves behind is what export_qdq writes and what qdq_stats reads back:

sample = torch.randn(64, 3, 64, 64)  # the batch size the graph was captured at

qdq_path = export_qdq(pt2e.model.cpu(), sample, TMP/"resnet18_qat_int8.onnx")
print(qdq_path.name, f"{qdq_path.stat().st_size / 1e6:.1f} MB")
print(qdq_stats(qdq_path).as_dict())

[torch.onnx] Obtain model graph for `GraphModule([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `GraphModule([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
resnet18_qat_int8.onnx 11.4 MB
{'n_quantize': 33, 'n_dequantize': 54, 'n_per_channel': 21, 'n_nonzero_zero_point': 0, 'n_unquantized_conv_add': 0}

33 QuantizeLinear and 54 DequantizeLinear nodes, 21 tensors with one scale per channel, 0 non-zero zero-points — what symmetric INT8 promises, counted on the file rather than trusted — and 0 Add inputs read straight from a convolution, which is the default qdq_placement='per_op'.

Measured on

ResNet-18 (ImageNet weights, fc replaced), Oxford-IIIT Pet as a two-class problem, 64px, batch 64, torch 2.9.1+cu128, single run, torch.manual_seed(0) set before each arm. Both accuracies score the model each arm produced — the floating-point model for the first arm, the converted INT8 model for the second — on the CPU over the whole validation set, n = 1478; the pt2e arm drops the six-image last batch and reports no accuracy. No latency is measured. The two sizes are state_dict bytes; the ONNX size is the file export_qdq wrote.

Summary

Tool	What it gives you
`QuantizeCallback(backend='x86')`	An FX INT8 model on `learn.model` after the fit, converted from the fake-quantize modules it trained through
`QuantizeCallback(backend='pt2e')`	A symmetric INT8 graph, refused up front when the dataloaders yield more than one batch size
`quant_spec(model)`	The precision record the callback attached to the model it produced
`export_qdq(model, sample, path)`	A QDQ ONNX file written from that graph
`qdq_stats(path)`	Node counts, per-channel count and zero-point audit, read off the file