torch.compile deep dive

torch.compile deep dive#

Track 07 - GPU · Notebook 06 · Runtime: ≈2 min on CPU or GPU

Prerequisites: 07_gpu/01 (GPU architecture tour).

References:

torch.compile documentation.

Ansel et al. 2024, PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation (PyTorch 2 paper).

What#

torch.compile(model) wraps model.__call__ in TorchDynamo, which traces the Python-level forward pass, emits a graph, and hands it to a backend (inductor by default) that fuses pointwise ops, picks memory layouts, and generates Triton kernels for the fused pieces.

Two things to internalise:

Graph breaks. Any dynamic control flow, untraceable Python call, or unknown external library forces Dynamo to split the forward into two graphs at that point and fall back to eager Python in between. Each break adds a Python re-entry (typically 50–500 µs) and blocks cross-op fusion across the seam — enough to halve the compile speedup on a small forward, and to wipe out the entire reduce-overhead mode benefit on a large one. The target is ≤ 1 graph break per forward, ideally zero.
Dynamic shapes. Every unique input shape triggers a recompile unless the model is compiled with dynamic=True or shapes are marked dynamic. Production LLM inference wants mode="reduce-overhead"
- dynamic=True to keep recompile latency out of steady-state.

We compile a small MLP three ways (eager, default compile, reduce-overhead compile), measure the compile+warmup time, then compare steady-state latency. Everything runs on CPU if no GPU is available.

from llm_systems_cookbook.nb import bootstrap

import time

import torch
import torch.nn as nn

s = bootstrap("07_gpu_06_torch_compile_deep_dive")

A small but non-trivial MLP#

Two hidden layers with SiLU activations. Big enough that kernel launch overhead is meaningful vs actual compute.

class MLP(nn.Module):
    def __init__(self, d: int = 1024, hidden: int = 4096) -> None:
        super().__init__()
        self.fc1 = nn.Linear(d, hidden)
        self.fc2 = nn.Linear(hidden, hidden)
        self.fc3 = nn.Linear(hidden, d)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc3(torch.nn.functional.silu(self.fc2(torch.nn.functional.silu(self.fc1(x)))))


torch.manual_seed(0)
model = MLP().to(DEVICE).eval()
params = sum(p.numel() for p in model.parameters())
print(f"model params = {params/1e6:.2f}M  on {DEVICE}")

Three variants#

eager: pristine PyTorch, no compilation.
compile_default: torch.compile(model), default mode="default".
compile_reduce_overhead: torch.compile(model, mode="reduce-overhead"). This mode uses CUDA graphs and aggressive caching; on CPU it falls back to "default" but the intent is to drop per-call dispatch overhead for short forward passes.

torch._dynamo.reset() between variants ensures each gets a clean cache so we measure compile time too.

import torch._dynamo as dynamo

def bench_forward(fn, x, *, warmup: int = 3, iters: int = 30) -> tuple[float, float]:
    '''Return (warmup_total_s, steady_iters_per_s).'''
    t0 = time.perf_counter()
    for _ in range(warmup):
        fn(x)
    if IS_CUDA:
        torch.cuda.synchronize()
    warm_s = time.perf_counter() - t0

    t0 = time.perf_counter()
    for _ in range(iters):
        fn(x)
    if IS_CUDA:
        torch.cuda.synchronize()
    return warm_s, iters / (time.perf_counter() - t0)


B = 16
x = torch.randn((B, 1024), device=DEVICE)

dynamo.reset()
eager_fn = model
warm_eager, its_eager = bench_forward(eager_fn, x)
print(f"eager              warm={warm_eager*1000:7.1f} ms   steady={its_eager:7.1f} it/s")

dynamo.reset()
compiled_default = torch.compile(model, mode="default")
warm_def, its_def = bench_forward(compiled_default, x)
print(f"compile(default)   warm={warm_def*1000:7.1f} ms   steady={its_def:7.1f} it/s")

dynamo.reset()
compiled_ro = torch.compile(model, mode="reduce-overhead")
warm_ro, its_ro = bench_forward(compiled_ro, x)
print(f"compile(reduce-ov) warm={warm_ro*1000:7.1f} ms   steady={its_ro:7.1f} it/s")

Graph break detection#

A graph break happens when Dynamo can’t trace a line - an un-traceable Python call, a print, a numpy interop, or dynamic shape code. torch._dynamo.explain inspects a function and reports how many graphs Dynamo produces; one graph = no break, two+ = broken.

def clean_forward(x: torch.Tensor) -> torch.Tensor:
    return model(x)


def broken_forward(x: torch.Tensor) -> torch.Tensor:
    y = model(x)
    # Force an explicit break so the check is version-stable.
    torch._dynamo.graph_break()
    return model(y)


explanation_clean = dynamo.explain(clean_forward)(x)
explanation_broken = dynamo.explain(broken_forward)(x)

n_graphs_clean = getattr(explanation_clean, "graph_count", None) or len(getattr(explanation_clean, "graphs", []))
n_graphs_broken = getattr(explanation_broken, "graph_count", None) or len(getattr(explanation_broken, "graphs", []))
print(f"clean forward:   graphs = {n_graphs_clean}  break_count = {getattr(explanation_clean, 'break_reasons', [])}")
print(f"broken forward:  graphs = {n_graphs_broken}  break_count = {getattr(explanation_broken, 'break_reasons', [])}")

s.check(
    "clean_forward_single_graph",
    lambda: n_graphs_clean <= 1,
    msg=f"graph_count = {n_graphs_clean}",
)
s.check(
    "broken_forward_has_more_graphs",
    lambda: n_graphs_broken > n_graphs_clean,
    msg=f"clean={n_graphs_clean}  broken={n_graphs_broken}",
)

Performance checks#

compile(default) steady-state ≥ eager (often a lot more; at minimum equal). On CPU the inductor backend can sometimes be slightly slower for tiny models due to overhead; we allow parity.
compile warmup is strictly larger than eager warmup (compilation costs real time - that’s the tradeoff).

s.check(
    "compile_default_at_least_matches_eager",
    lambda: its_def >= 0.90 * its_eager,
    msg=f"eager {its_eager:.1f} it/s  compile(default) {its_def:.1f} it/s",
)
s.check(
    "compile_warmup_exceeds_eager_warmup",
    lambda: warm_def >= warm_eager,
    msg=f"warm_eager={warm_eager*1000:.1f} ms  warm_compile={warm_def*1000:.1f} ms",
)

Cold vs warm latency per mode#

The compile-time numbers and the steady-state numbers tell different stories. Plot both together: the left bar is the first-call cost (includes Dynamo trace + Inductor codegen + kernel cache warm-up); the right bar is per-iteration at steady state. The tradeoff lives in the ratio between them.

import matplotlib.pyplot as plt

modes   = ["eager", "compile (default)", "compile (reduce-ov)"]
warm_ms = [warm_eager * 1000.0, warm_def * 1000.0, warm_ro * 1000.0]
steady_ms = [1000.0 / its_eager, 1000.0 / its_def, 1000.0 / its_ro]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.3))
colors = ["tab:gray", "tab:blue", "tab:green"]
ax1.bar(modes, warm_ms, color=colors)
ax1.set_ylabel("cold / warmup time (ms)")
ax1.set_title("first call: compile + trace cost")
for t in ax1.get_xticklabels():
    t.set_rotation(15)
ax1.grid(True, axis="y", alpha=0.3)

ax2.bar(modes, steady_ms, color=colors)
ax2.set_ylabel("steady-state latency (ms / iter)")
ax2.set_title("warm call: amortised over many iters")
for t in ax2.get_xticklabels():
    t.set_rotation(15)
ax2.grid(True, axis="y", alpha=0.3)

fig.suptitle(f"torch.compile: cold vs warm  ({DEVICE})")
fig.tight_layout()
plt.show()

Exercises#

Enable dynamic=True and rerun with batches of different sizes. Without it, each new shape triggers a recompile; with it, compile once and reuse.
Inspect inductor IR: set TORCH_COMPILE_DEBUG=1 and rerun. Look at the generated Triton kernels in torch_compile_debug/….
Profile the graph breaks: the broken_forward above has two intentional breaks. Identify both using dynamo.explain and rewrite the function to trace in one graph.

References#

PyTorch 2 paper for the architecture (Dynamo, AOTAutograd, Inductor).
torch._dynamo.explain and torch._dynamo.config.verbose for debugging compile failures.
torchtune for examples of torch.compile applied to LLM training with dynamic=True.

s.summary()
s.save()