TorchDynamo: capturing Python at runtime

Because it fails silently and plausibly. A limitation you hit is an error message; a limitation that produces wrong numbers is a bug that ships.

def forward(self, x):
    if x.sum() > 0:        # evaluated once, during tracing
        return self.a(x)
    return self.b(x)       # if the example was positive, this branch is gone forever

The traced module contains only self.a. It runs, returns tensors of the right shape and dtype, and is wrong for half its inputs. Nothing warns you. The same applies to for i in range(x.shape[0]), which unrolls to the example’s batch size and then produces garbage or a shape error later, and to any Python int derived from a tensor, which is frozen as a constant.

Contrast the two safe designs. script refuses to compile what it cannot represent — correct, but it moved the burden to the user. Dynamo represents what it can and falls back to the interpreter for the rest, so the failure mode is a performance loss rather than a correctness loss.

That ordering is the design principle worth extracting: when a capture system meets something it cannot express, its options are be wrong, refuse, or deoptimise. Only the third gives you both correctness and adoption, and it is why every successful JIT for a dynamic language has chosen it.

Take a break in the middle of a function. Dynamo compiles the ops it captured before the break into graph A, emits a resume function whose bytecode continues from the break point, runs the offending operation in the ordinary Python interpreter, then re-enters compilation for graph B.

The visible cost is modest: two compiled graphs instead of one, plus a Python round trip.

The costs that are not visible are the ones that matter:

• No fusion across the boundary. The last op of A and the first of B cannot be fused, so an intermediate that would have stayed in registers is now materialised to HBM and read back — the Module 3 traffic argument, in reverse.
• CUDA Graphs are dead across it. A replayable graph cannot contain a Python callback, so mode="reduce-overhead" degrades to capturing each fragment separately, or to nothing. For a launch-bound model this alone can erase the entire benefit (Module 9).
• Guards multiply. Each fragment carries its own guard set, checked on every call.
• Breaks in loops compound. One break inside a per-layer loop is 32 breaks in a 32-layer model, so 33 graphs, none of which can fuse with its neighbours.

Which is why TORCH_LOGS=graph_breaks is the first thing to run when compilation underdelivers, and why fullgraph=True during development is worth the friction. A model reported as “compiled” with forty breaks has most of eager’s overhead and all of compilation’s cost.

A guard is a cheap predicate over the inputs and the environment that must hold for the compiled code to be valid. Typical contents: tensor dtype, device, layout, requires-grad, rank and shape (or a symbolic relation on shape, Module 10); the identity or value of Python arguments; the type of self and values of attributes it read; global state such as the grad mode and autocast state.

The design tension is a direct trade:

• Too strict and you recompile constantly. Guarding on an exact batch size recompiles for every batch size; hit cache_size_limit (8) and Dynamo stops trying and runs eager forever — often the real reason a model “stopped being fast”.
• Too loose and you run a compiled graph under assumptions that no longer hold, which is miscompilation: silently wrong numbers.

So guards must be simultaneously sound (never admit an invalid input), precise (not so conservative as to recompile needlessly), and cheap (a few microseconds, since they run on every call and the whole point was to remove microseconds).

The clever move is symbolic shapes: rather than guarding batch == 4, guard batch >= 2 or record that two dimensions are equal, so one compiled artifact serves many shapes. That is Module 10’s subject and it is where most remaining pain in torch.compile lives — because deciding which facts to specialise on and which to leave symbolic is undecidable in general, and the heuristics are what you fight.

All three have been tried; bytecode is the only one that is both correct and transparent.

Source-level (AutoGraph in TF2, and jit.script’s parser) means parsing Python, rewriting control flow, and handling every syntactic form the language allows — including code you did not write, in libraries you do not control. You end up debugging generated source, and you must reimplement Python’s semantics correctly, which nobody has ever quite done.

Op-level tracing (jit.trace, fx.symbolic_trace) sees only what executed, so control flow is invisible by construction. That is the silent-wrongness failure above; you cannot fix it by trying harder, because the information was never observed.

Bytecode sits at exactly the right level. It is a small, stable, well-specified instruction set — far simpler than Python’s surface syntax. Control flow is explicit as jump instructions, so Dynamo can see a branch rather than merely following one. It works regardless of how the source was written, including decorated, generated or third-party code. And because PEP 523 lets you replace the evaluation of a frame, Dynamo can hand back modified bytecode that calls the compiled graph, so the fallback path is the interpreter itself rather than a reimplementation of it.

The cost is that Dynamo must model the CPython interpreter’s stack machine, and is coupled to bytecode details that change between Python versions — which is exactly why each new CPython release needs Dynamo work. That is a maintenance burden accepted deliberately in exchange for being able to capture arbitrary Python correctly.