PyTorch Dynamic Quantization

Dynamic quantization stores Linear weights as int8 numbers in a converted model. An int8 number is a whole number that fits in one byte.

One layer, two sizes

A wide Linear layer holds thousands of float weights. The wrapper here is a one-layer Sequential, and its quantized form stores each weight in one byte.

import io
import torch

torch.manual_seed(0)
model = torch.nn.Sequential(torch.nn.Linear(64, 64))

def size_of(m):
    buf = io.BytesIO()
    torch.save(m.state_dict(), buf)
    return buf.tell()

quantized = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
print(f"Plain: {size_of(model)} bytes")
print(f"Quantized: {size_of(quantized)} bytes")

The weight dtype changes

The plain weight is float32, and the converted layer serves its weight as qint8. The check asks each layer for its stored weight.

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(torch.nn.Linear(4, 3))
quantized = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
print(model[0].weight.dtype)
print(quantized[0].weight().dtype)

A quantized model still runs

The converted model answers a small batch right away. Its outputs sit close to the plain outputs, and every value moves a little.

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(torch.nn.Linear(4, 3))
quantized = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
x = torch.tensor([[1.0, 0.0, 0.0, 0.0]])
with torch.no_grad():
    plain = model(x)
    got = quantized(x)
print(f"Plain: {[round(v, 4) for v in plain[0].tolist()]}")
print(f"Quantized: {[round(v, 4) for v in got[0].tolist()]}")

Example

import io
import torch

CENTERS = torch.randn(4, 32, generator=torch.Generator().manual_seed(0)) * 2.0

def make(rows, seed):
    g = torch.Generator().manual_seed(seed)
    labels = torch.randint(0, 4, (rows,), generator=g)
    data = CENTERS[labels] + torch.randn(rows, 32, generator=g) * 0.5
    return data, labels

train_x, train_y = make(240, 1)
test_x, test_y = make(200, 2)

torch.manual_seed(3)
model = torch.nn.Sequential(
    torch.nn.Linear(32, 128), torch.nn.ReLU(), torch.nn.Linear(128, 4))
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
loss_fn = torch.nn.CrossEntropyLoss()
for step in range(150):
    optimizer.zero_grad()
    loss = loss_fn(model(train_x), train_y)
    loss.backward()
    optimizer.step()

def size_of(m):
    buf = io.BytesIO()
    torch.save(m.state_dict(), buf)
    return buf.tell()

def predictions(m, x):
    with torch.no_grad():
        return m(x).argmax(dim=1)

model.eval()
plain_pred = predictions(model, test_x)
quantized = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8)
quantized.eval()
quant_pred = predictions(quantized, test_x)
plain_acc = (plain_pred == test_y).float().mean().item()
quant_acc = (quant_pred == test_y).float().mean().item()
agree = (plain_pred == quant_pred).float().mean().item()
print(f"Plain size: {size_of(model)} bytes")
print(f"Quantized size: {size_of(quantized)} bytes")
print(f"Plain accuracy: {plain_acc:.2f}")
print(f"Quantized accuracy: {quant_acc:.2f}")
print(f"Top-1 agreement: {agree:.2f}")

Key points

Qint8: the name of the one-byte integer dtype for quantized values
Activations: the values that flow between layers while a model runs
Use quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8) to convert a model
Dynamic means activations are converted on the fly during each run
Only the named layer types convert, and other layers stay float
Rounded weights shift predictions a little, and the shift can be counted

In this example, a two-layer classifier trains on four seeded blob classes. The data has 32 features, and the hidden layer holds 128 units. Sizes come from saving each state_dict into memory and counting bytes. The plain model prints 1.00 accuracy on the held-out batch. The quantized model also prints 1.00 accuracy on the same batch. The agreement between their top answers prints as 1.00.

Note

Torch 2.14 warns that eager dynamic quantization is a legacy path

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import io
import torch

CENTERS = torch.randn(4, 32, generator=torch.Generator().manual_seed(0)) * 2.0

def make(rows, seed):
    g = torch.Generator().manual_seed(seed)
    labels = torch.randint(0, 4, (rows,), generator=g)
    data = CENTERS[labels] + torch.randn(rows, 32, generator=g) * 0.5
    return data, labels

train_x, train_y = make(240, 1)
test_x, test_y = make(200, 2)

torch.manual_seed(3)
model = torch.nn.Sequential(
    torch.nn.Linear(32, 128), torch.nn.ReLU(), torch.nn.Linear(128, 4))
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
loss_fn = torch.nn.CrossEntropyLoss()
for step in range(150):
    optimizer.zero_grad()
    loss = loss_fn(model(train_x), train_y)
    loss.backward()
    optimizer.step()

def size_of(m):
    buf = io.BytesIO()
    torch.save(m.state_dict(), buf)
    return buf.tell()

def predictions(m, x):
    with torch.no_grad():
        return m(x).argmax(dim=1)

model.eval()
plain_pred = predictions(model, test_x)
quantized = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8)
quantized.eval()
quant_pred = predictions(quantized, test_x)
plain_acc = (plain_pred == test_y).float().mean().item()
quant_acc = (quant_pred == test_y).float().mean().item()
agree = (plain_pred == quant_pred).float().mean().item()
print(f"Plain size: {size_of(model)} bytes")
print(f"Quantized size: {size_of(quantized)} bytes")
print(f"Plain accuracy: {plain_acc:.2f}")
print(f"Quantized accuracy: {quant_acc:.2f}")
print(f"Top-1 agreement: {agree:.2f}")

Exercise

What does dynamic quantization store as int8?

The Linear weights The training rows The labels The gradients