PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation

PyTorch runs Python line by line, which is easy to debug but slow, so torch.compile rewrites the program's bytecode while it runs to capture the maths into a graph, then compiles that graph into fast GPU code.

Why the flexible version is slow

When you write PyTorch, each line does its work immediately:

y = a + b        # runs now, produces a tensor
z = y.relu()     # runs now, produces another tensor
w = z * 2        # runs now

This is called eager mode — the framework is eager, and does each operation the moment it is asked. The advantage is enormous for a person writing models. You can print a value in the middle. You can set a breakpoint. An error points at the line that caused it.

The cost is hidden in the memory system. Each of those three lines writes a complete tensor to GPU memory, and the next line reads it straight back. For large tensors, the arithmetic is trivial and the memory traffic dominates. The GPU spends most of its time moving numbers rather than computing with them.

A compiler that could see all three lines at once would do something obvious: compute (a + b), apply relu, and multiply by two while the values are still in registers, writing to memory only once. Combining several operations into one like this is called fusion. Eager mode cannot do it, because when line one runs, lines two and three have not happened yet.

What people tried before, and where it broke

The alternative is graph mode: capture the whole computation as a graph first, then hand that graph to a compiler.

The difficulty is the capture step, and the reason is that PyTorch models are written in Python — a language where almost anything can happen at any point. A model’s forward method may contain a conditional that depends on the input values, a loop whose length depends on the data, a call into a third-party library, or a print.

Earlier capture approaches ran into one of three walls:

• Tracing runs the model once with example inputs and records what happened. Anything data-dependent is silently baked in. A conditional that took the if branch during tracing will take it forever, even when it should not.
• Source parsing reads the Python text and tries to translate it. It fails on anything dynamic, and Python is a very dynamic language.
• Rewrite the model puts the burden on the user, which defeats the purpose of a framework people chose for being pleasant to write.

The abstract states the goal precisely: capture graphs without sacrificing the flexibility of Python.

The idea: intercept below Python, not above it

Both earlier approaches worked at the level of the Python source. TorchDynamo works one level lower, at the level of Python bytecode.

Python does not execute your source text directly. It first compiles each function into

This position is what makes the difference. At the bytecode level, TorchDynamo can see exactly what the interpreter is about to do, including the actual types and values involved, which the source text alone never reveals.

It then divides the function into two kinds of work:

1. Operations it understands — PyTorch tensor operations — which it extracts into an FX graph, PyTorch’s own graph representation.
2. Everything else — the prints, the unusual library calls, the parts it cannot reason about — which it leaves running as normal Python.

This division is the key move. Earlier systems treated an unsupported construct as a failure of the whole capture. TorchDynamo treats it as a boundary: capture what you can, run the rest normally, and carry on.

model = torch.compile(model)
# no work done yet

Why it works

Because interception happens at the bytecode level rather than the source level, TorchDynamo sees the actual operations and actual types about to execute. → Because it sees them, it can decide operation by operation what it understands, instead of accepting or rejecting the whole function. → Because unsupported code becomes a boundary rather than a failure, capture succeeds on real models that contain prints, conditionals and third-party calls. → Because the captured graph contains several operations rather than one, the backend can fuse them and write to memory once instead of after every step. → Because compiled results are reused across calls, the compilation cost is paid once and the speedup accrues on every later call.

That chain is what the abstract summarises as capturing graphs more robustly than prior approaches while adding minimal overhead.

What it costs

• The first call is slow. Capture and compilation happen at run time. A program that calls a function only a few times may never earn the cost back.
• Two new systems to understand. When results are wrong or slow, the engineer now debugs a capture stage and a compilation stage as well as their model.
• The escape hatch has a price. Falling back to Python keeps capture working, but each fallback shrinks the graph, and a smaller graph gives the compiler less to optimise across.

What the measurements showed

The abstract reports the following, and this page reports nothing beyond it.