Inference Optimization — torch.compile, Quantization, Batching

Training is iterative; inference is performance-critical

Once the model trains, you optimize it for serving. Three techniques you'll combine routinely:

torch.compile(model) — JIT-compile the model graph with TorchInductor. 1.5-3x speedup with one line of code. (Full coverage in the next track.)
Quantization — reduce numerical precision (fp32 → int8 or int4) for smaller model and faster matmul. Modern path is torchao.
Batching — instead of running 32 single-sample inferences, run one 32-sample batch. Underutilization is real cost.

The order to apply them

Get correctness right with eager-mode + fp32. Verify outputs.
Switch to bf16 if your hardware supports it. Verify accuracy holds.
Add torch.compile(model). Verify speed and accuracy.
Quantize if you need more (int8 dynamic for transformers, int8/int4 weight-only for LLMs).
If you're serving requests, batch across users with a small queueing window.

Don't forget the Python overhead

Tokenization, post-processing, serialization can dominate inference time on small models. Profile end-to-end (HTTP request → response), not just the forward pass. The fix is often "cache the tokenizer" or "use a faster JSON library" rather than anything model-side.

Code

torch.compile — the one-line speedup·python

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english",
)
model.eval()

# JIT-compile. First call is slow (compilation), subsequent calls are fast.
model = torch.compile(model)

tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
inputs = tok(["A great movie!", "Terrible."], padding=True, return_tensors="pt")

with torch.inference_mode():
    out = model(**inputs)

print(out.logits.argmax(-1))   # tensor([1, 0])

Dynamic int8 quantization — the one-liner for transformers·python

import torch
import torch.nn as nn
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english"
)
model.eval()

# Quantize Linear layers to int8 — weights stored as int8, computed in int8
quantized = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8,
)

# Compare disk size
torch.save(model.state_dict(), '/tmp/fp32.pt')
torch.save(quantized.state_dict(), '/tmp/int8.pt')

import os
print(f"fp32 : {os.path.getsize('/tmp/fp32.pt')/1e6:.1f} MB")
print(f"int8 : {os.path.getsize('/tmp/int8.pt')/1e6:.1f} MB")
# fp32: 268.4 MB
# int8:  72.1 MB     (~4x smaller)

torchao — modern int4 / weight-only quantization·python

# pip install torchao
import torchao
from torchao.quantization import int4_weight_only, int8_weight_only
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("gpt2")

# int8 weight-only — weights stored as int8, computed in fp16/bf16
torchao.quantize_(model, int8_weight_only())

# int4 weight-only — even smaller, designed for LLMs
# torchao.quantize_(model, int4_weight_only(group_size=128))

# Use the model normally — the int4/int8 layers handle on-the-fly dequant
# in their forward.

End-to-end inference benchmark·python

import time
import torch

@torch.inference_mode()
def benchmark(model, inputs, n=100, warmup=10):
    if next(model.parameters()).is_cuda:
        torch.cuda.synchronize()
    for _ in range(warmup):
        model(**inputs)
    if next(model.parameters()).is_cuda:
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n):
        model(**inputs)
    if next(model.parameters()).is_cuda:
        torch.cuda.synchronize()
    elapsed_ms = (time.perf_counter() - t0) / n * 1000
    print(f"{elapsed_ms:.2f} ms / call")
    return elapsed_ms

Inference Optimization — torch.compile, Quantization, Batching

Training is iterative; inference is performance-critical

The order to apply them

Don't forget the Python overhead

Code

External links

Exercise

Progress

Comments 0