Compute Graphs and Compilation

The Graph Was Always There

Autograd already builds a compute graph — to run backward.

Eager: walk it one kernel at a time (a launch + a memory round trip each). Capture it, and a compiler can rewrite across nodes. That is the whole idea.

Two Capture Philosophies

torch.compile: capture Python bytecode, graph-break to Python when stuck. jax.jit: trace to a jaxpr — no breaks, but recompile on new shapes.

Breaks vs. Retraces

def f(x):
    x = x * 2
    if x.sum() > 0:        # Data-dependent: forces a graph break
        return x + 1
    return x - 1

explanation = torch._dynamo.explain(f)(torch.randn(8, device=d2l.try_gpu()))
print(f'graph breaks: {explanation.graph_break_count}')

graph breaks: 1

@jax.jit
def f(x):
    print('this prints once, at trace time')  # Not in the compiled graph
    return jnp.sin(x) + 1

x = jnp.arange(4.0)
print(jax.make_jaxpr(lambda x: jnp.sin(x) + 1)(x))  # The captured graph
_ = f(x); _ = f(x)   # Second call: no print — the trace is cached

The print vanishes: tracing sees tensor ops only. Purity is the price of having no graph breaks.

Fusion Cures the Bandwidth Chain

The bleeding chain from §13.1, cured in one line:

x = torch.randn(4000, 4000, device=d2l.try_gpu())

def gelu_ish(x):
    return 0.5 * x * (1 + torch.tanh(0.8 * (x + 0.04 * x**3)))

compiled = torch.compile(gelu_ish)
# First call compiles; then verify the rewrite: same answer, then faster
assert torch.allclose(gelu_ish(x), compiled(x), atol=1e-6)

print(d2l.Benchmark(lambda: gelu_ish(x), desc='eager'))
print(d2l.Benchmark(lambda: compiled(x), desc='compiled'))

eager: 1.25 ms/call
compiled: 0.15 ms/call

Same answer first (the allclose), then faster: close to an order of magnitude, entirely on the bytes side — one memory round trip instead of one per op. FlashAttention (§10.5) is this idea by hand.

Compile the Training Step

class GeluIsh(nn.Module):  # Wraps the elementwise chain above as a layer
    def forward(self, x):
        return gelu_ish(x)

net = nn.Sequential(nn.Linear(1024, 1024), GeluIsh(),
                    nn.Linear(1024, 1024), GeluIsh(),
                    nn.Linear(1024, 1024), GeluIsh(),
                    nn.Linear(1024, 1024)).to(d2l.try_gpu())
X = torch.randn(512, 1024, device=d2l.try_gpu())
opt = torch.optim.SGD(net.parameters(), lr=0.01)

def train_step(model):
    opt.zero_grad(set_to_none=True)
    model(X).sum().backward()
    opt.step()

cnet = torch.compile(net)
t0 = time.perf_counter()
train_step(cnet); torch.cuda.synchronize()  # First call: compiles
print(f'first compiled step: {time.perf_counter() - t0:.1f} s')

print(d2l.Benchmark(lambda: train_step(net), desc='eager'))
print(d2l.Benchmark(lambda: train_step(cnet), desc='compiled'))

first compiled step: 0.3 s
eager: 1.84 ms/call
compiled: 1.28 ms/call

Fixed cost on call one; repaid over every step after. torch.compile(net) captures forward+backward (the optimizer stays eager); jax.jit swallows the whole step, update included. JAX’s AOT lower().compile() also lets you inspect the plan.

The Overhead Regime: Capture & Replay

Deep stack of thin layers — launches dominate arithmetic:

# Many small layers: launch overhead dominates the actual arithmetic
deep = nn.Sequential(*[m for _ in range(60)
                       for m in (nn.Linear(256, 256), nn.Tanh())]).to(
    d2l.try_gpu())
x = torch.randn(64, 256, device=d2l.try_gpu())

reduced = torch.compile(deep, mode='reduce-overhead')
with torch.no_grad():  # Forward only: replay wants fixed buffers
    reduced(x)  # Warmup: compiles and captures the CUDA graph
    print(d2l.Benchmark(lambda: deep(x), desc='eager'))
    print(d2l.Benchmark(lambda: reduced(x), desc='reduce-overhead'))

eager: 2.33 ms/call
reduce-overhead: 0.31 ms/call

reduce-overhead records the launches into a CUDA graph and replays them as one. XLA amortizes launches by construction — no such knob needed.

When Not to Compile

• short runs — compile cost never repaid
• graph breaks in the hot loop — little captured
• shape churn (JAX) — recompiles every step
• already compute-bound — nothing to fuse

Diagnose first (§13.1). Compile for bandwidth and overhead, not reflexively.