Pruning, Knowledge Distillation, and Benchmarking

Three more techniques to make production models smaller and faster

Pruning

Pruning zeros out weights judged unimportant (by magnitude or by a structured rule). Two flavors:

Unstructured — zero individual weights. Theoretically reduces flops but most hardware doesn't speed up sparse matmul, so the practical wins are modest.
Structured — remove whole channels / heads / blocks. Fewer params AND fewer flops AND smaller activations. Where the practical wins are.

Pruning + quantization compose well. Pruning + distillation (train a smaller student from the pruned teacher) compose even better.

Knowledge distillation

Train a small "student" model to match the soft predictions of a large "teacher" model. The student learns from the teacher's logits, not just the hard labels — which carry more information ("this is mostly cat but a little fox"). The standard recipe:

Run teacher on the input — get soft logits.
Run student on the input — get its logits.
Loss = α · KL(student / T || teacher / T) · T² + (1 − α) · CE(student, hard_label)

where T is the "temperature" (typically 2-8) that softens both distributions, and α weights soft vs hard targets.

Benchmarking — your honest report card

Every optimization claim deserves a benchmark. Measure latency at the actual batch size you'll deploy with, on the actual hardware, with the actual input distribution. Anecdotal speedups don't survive contact with production.

Code

Unstructured pruning — magnitude-based·python

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

linear = nn.Linear(100, 50)

# Zero out 30% of weights by magnitude (smallest go to zero)
prune.l1_unstructured(linear, name='weight', amount=0.3)

# Check sparsity
sparsity = (linear.weight == 0).float().mean().item()
print(f"sparsity: {sparsity:.2%}")           # ~30%

# Make pruning permanent (removes the mask, keeps the zeros)
prune.remove(linear, 'weight')

Knowledge distillation — the classic recipe·python

import torch
import torch.nn as nn
import torch.nn.functional as F

teacher = BigModel().eval()                   # frozen
student = SmallModel().train()
optimizer = torch.optim.AdamW(student.parameters(), lr=1e-3)
T = 4.0                                        # temperature
alpha = 0.7                                    # soft-vs-hard weight

for x, y in train_loader:
    optimizer.zero_grad()

    with torch.inference_mode():
        teacher_logits = teacher(x)

    student_logits = student(x)

    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.softmax(teacher_logits / T, dim=-1),
        reduction='batchmean',
    ) * (T ** 2)

    hard = F.cross_entropy(student_logits, y)

    loss = alpha * soft + (1 - alpha) * hard
    loss.backward()
    optimizer.step()

Benchmark utility — the table everyone needs·python

import time
import torch

@torch.inference_mode()
def benchmark(model, input_tensor, num_runs=100, warmup=10, label='model'):
    is_cuda = next(model.parameters()).is_cuda
    if is_cuda: torch.cuda.synchronize()

    for _ in range(warmup):
        model(input_tensor)
    if is_cuda: torch.cuda.synchronize()

    t0 = time.perf_counter()
    for _ in range(num_runs):
        model(input_tensor)
    if is_cuda: torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0

    avg_ms = elapsed / num_runs * 1000
    throughput = input_tensor.size(0) * num_runs / elapsed
    print(f"{label:20s} avg {avg_ms:6.2f} ms/call, {throughput:,.0f} samples/sec")
    return avg_ms

# Use it like:
benchmark(fp32_model, batch, label='fp32 baseline')
benchmark(quant_model, batch, label='int8 quant')
benchmark(compiled_quant_model, batch, label='int8 + compile')

Exercise

Implement the benchmark utility from the third code block. Use it on three configurations of any small classifier: (a) fp32 eager, (b) bf16 + compile, (c) int8 dynamic quant. Print the table. The size of the latency wins (or the flatness of them) tells you whether the model is worth optimizing further.