🎓 A1Learner
🔍
0 of 71 lessons completed

PyTorch torch.compile on CPU

torch.compile is a tool that rewrites a model into optimized machine code before running it. The rewrite can join several small operations into one, a step called fusing.

Creating the wrapper compiles nothing yet

torch.compile returns a wrapper at once. The real compilation waits for the first call.

import torch

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

compiled = torch.compile(model)
print("Wrapper is callable:", callable(compiled))

The eager original, for comparison

The plain model runs at once with no compile step. Its output is the reference the compiled model must match.

x = torch.randn(32, 64)
with torch.no_grad():
    out = model(x)
print(out.shape)

The first call pays the compile wait

The first call triggers the rewrite and runs far longer than later calls. Later calls reuse the compiled code, so this snippet carries no stated output.

with torch.no_grad():
    first = compiled(x)
    second = compiled(x)
print(torch.equal(first, second))

Example

import time

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Linear(64, 64),
    torch.nn.ReLU(),
    torch.nn.Linear(64, 10),
)
model.eval()
x = torch.randn(32, 64)

start = time.perf_counter()
compiled = torch.compile(model)
with torch.no_grad():
    out_compiled = compiled(x)
compile_s = time.perf_counter() - start

with torch.no_grad():
    out_eager = model(x)

print("Output shape:", out_compiled.shape)
print("Outputs equal:", torch.equal(out_eager, out_compiled))
print(f"First compiled call took: {compile_s:.0f} s")

Key points

  • Compiled wrapper: the callable object that torch.compile returns for a model
  • Use torch.compile(model) to wrap a model for compilation
  • Compilation happens on the first call, not when the wrapper is created
  • Eager: the normal mode where each operation runs as Python reaches it
  • A compiled model must return the same outputs as the eager original
  • On this 2-core machine the first call for a tiny model took over a minute
  • At tiny model sizes the compiled runs showed no clear speed gain in measurement
  • In short, torch.compile trades a long first-call wait for a compiled model
  • This lesson uses torch 2.14.1+cpu, and compile behaviour changes between versions

In this example, a tiny three-layer model is wrapped with torch.compile. The clock starts before the first call, so the printed time is the compile wait plus one run. That call took 85 seconds on the recorded run, and a second run took 124 seconds. The eager original then produces its output with no wait at all. The outputs compare equal, which is the part worth trusting. No run-time speedup is claimed, because measurements at this tiny size showed none. On this class of machine, the compile wait is the price, and it is paid before any benefit can start.

Note

The closing example verifies correctness only, because the compile wait alone passes the usual time budget on this machine.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

When does torch.compile do its compilation work?