02 · `torch.compile` ⚡

The headline feature of PyTorch 2.x. One line wraps your model, and it often runs meaningfully faster — training and inference alike — with no change to your code.

---

1. The 'Why'

By default PyTorch runs in eager mode: each operation in your forward executes immediately, one at a time, as the Python interpreter reaches it. This is wonderful for development — it's why debugging PyTorch feels like debugging normal Python (Module 02) — but it leaves performance on the table. Every operation is dispatched separately, launching a fresh GPU kernel and shuffling intermediate results back and forth to memory, even when neighboring operations could be fused into one efficient kernel. For years, getting that performance meant exporting to a static graph (TorchScript) and accepting a clunkier workflow. PyTorch 2.x changed the deal.

torch.compile gives you graph-level speed without giving up eager-mode ergonomics. You wrap your model in a single call, keep writing and debugging normal PyTorch, and behind the scenes PyTorch traces your model into a graph, fuses operations, and generates optimized kernels — recompiling automatically when needed. The result is often a substantial speedup on modern GPUs, for both training loops and inference, with essentially zero code changes. It's not literally free (there's an upfront compilation cost, and some models benefit more than others), but it's the closest thing to a "make it faster" button that deep learning has, and understanding it is essential for anyone deploying PyTorch models today.

---

2. Core Concepts

What it does. torch.compile(model) returns an optimized version of your model. On the first call it captures and compiles the computation (fusing ops, generating fast kernels via a backend like TorchInductor); subsequent calls run the compiled version.

Drop-in and eager-compatible. The compiled model is used exactly like the original — same inputs, same outputs. You keep full eager-mode debuggability; compilation is transparent.

Works for training and inference. Unlike export (which targets deployment), torch.compile speeds up your ordinary training loop too — just compile the model before the loop.

Warmup / compilation cost. The first call (and any call with a new input shape) triggers compilation, which takes time. Speedups show up on the many subsequent calls, so it pays off over a full training run or high-volume serving, not a single prediction.

Recompilation on shape changes. Varying input shapes can trigger recompilation. Stable shapes (or dynamic=True) keep things fast.

---

3. Code in Action

The one-line speedup

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(1024, 1024), nn.ReLU(), nn.Linear(1024, 10))

compiled = torch.compile(model)      # <-- that's it; returns an optimized model
x = torch.randn(64, 1024)
out = compiled(x)                    # first call compiles; later calls run fast
print(out.shape)                     # torch.Size([64, 10])

Using it in a training loop

import torch
import torch.nn as nn

model = nn.Linear(1024, 10)
model = torch.compile(model)         # compile ONCE, before training
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for X, y in train_loader:            # the loop is unchanged from Module 05
    optimizer.zero_grad()
    loss = criterion(model(X), y)    # runs the compiled model
    loss.backward()
    optimizer.step()
# Early iterations pay the compile cost; the rest run faster.

Modes and dynamic shapes

import torch

# `mode` trades compile time for runtime speed
fast = torch.compile(model, mode="max-autotune")  # longer compile, best runtime

# If your input shapes vary (e.g. NLP sequence lengths), allow dynamic shapes
flexible = torch.compile(model, dynamic=True)      # avoids repeated recompilation

---

4. Common Pitfalls

Pitfall 1 — Benchmarking including the warmup call. The first forward pass triggers compilation and is slow; if you time it, you'll wrongly conclude torch.compile made things worse. Run a few warmup iterations first, then measure the steady-state — that's where the speedup lives.

Pitfall 2 — Expecting gains on tiny models or a single prediction. Compilation overhead only pays off across many calls. For a one-off inference on a small model, plain eager mode may be faster overall. Reach for torch.compile on real training runs or sustained high-volume serving.

Pitfall 3 — Constantly changing input shapes. Feeding a new shape every call can trigger repeated recompilation, erasing the benefit. Keep shapes stable (pad to fixed sizes) or pass dynamic=True so PyTorch compiles a shape-flexible version once.

---

5. Further Reading & Watch List

• 📖 Deep Learning with PyTorch + the PyTorch 2.x docs
• 🌐 PyTorch — "Introduction to torch.compile"
• 🛠️ torch.compile docs · PyTorch 2.x overview
• ▶️ PyTorch Conference talks on torch.compile internals

---

⬅️ Prev: 01 · Exporting Models · ➡️ Next: 03 · Inference Optimization