PyTorch Compiler Debugging: Graph Breaks, Guards, Recompiles and torch.compile

Explain this chapter with AI

Copy this prompt into ChatGPT, Claude, Gemini, a local model, or another AI.

Apply this chapter with AI

Copy this prompt into ChatGPT, Claude, Gemini, a local model, or another AI.

PyTorch: Zero to Hero โ€” Advanced Step 09A

In the previous chapter we treated torch.compile as one tool inside a larger performance-debugging workflow.

That is enough until compilation itself becomes the problem.

Then the questions change.

The model runs in eager mode.

The model may even run when compiled.

But perhaps:

compile time is enormous
first call takes seconds
new shapes keep compiling again
graph breaks appear inside forward()
compiled execution is slower than eager
one backend works and another fails
small code changes produce very different compiler behavior

At that point, torch.compile(model) is not the end of the story.

It is the beginning of a different kind of debugging.

The central idea of this chapter is:

A compiled PyTorch program is still your Python program, but PyTorch must make assumptions about which parts can become graphs and when those graphs remain valid.

If we can inspect those assumptions, compiler debugging becomes much less mysterious.


1. Start with the eager program

Before debugging compilation, prove the ordinary program works.

import torch
import torch.nn as nn


class TinyMLP(nn.Module):
    def __init__(self, dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, dim * 2),
            nn.GELU(),
            nn.Linear(dim * 2, dim),
        )

    def forward(self, x):
        return self.net(x)


model = TinyMLP().eval()
x = torch.randn(64, 256)

with torch.inference_mode():
    eager_output = model(x)

print(eager_output.shape)

Now compile the same model:

compiled = torch.compile(model)

with torch.inference_mode():
    compiled_output = compiled(x)

And compare correctness before speed:

torch.testing.assert_close(
    eager_output,
    compiled_output,
    rtol=1e-4,
    atol=1e-5,
)

This gives us the first compiler-debugging split:

fails eager
    โ†’ model / PyTorch program problem

works eager, fails compiled
    โ†’ compiler-path problem

That single distinction saves a great deal of wasted time.


2. What torch.compile is trying to do

A useful simplified picture is:

    flowchart LR
    A[Python program] --> B[TorchDynamo]
    B --> C[FX graph]
    C --> D[AOTAutograd]
    D --> E[TorchInductor]
    E --> F[Generated optimized code]
  

The exact implementation is richer than this.

But the layers are useful because failures can happen at different places.

Very roughly:

TorchDynamo
    captures Python execution into graphs

AOTAutograd
    transforms forward/backward computation

TorchInductor
    lowers graphs toward optimized kernels/code

That gives us a debugging principle:

Do not ask only whether torch.compile failed. Ask which layer first stopped behaving as expected.


3. The most important word: graph

Suppose we write:

def fn(x):
    a = torch.sin(x)
    b = torch.cos(a)
    return b * 2

PyTorch can reason about the tensor operations as a graph:

x
โ†“
sin
โ†“
cos
โ†“
ร— 2
โ†“
output

Graph form is valuable because the compiler can see a region of computation rather than executing every Python statement independently.

Once the compiler has a graph, it can attempt optimizations across the region.

But Python can do almost anything.

That is where graph breaks enter the story.


4. A graph break means PyTorch stopped capturing

Consider:

@torch.compile
def fn(x):
    y = torch.sin(x)

    if x.sum().item() > 0:
        y = y * 2

    return torch.cos(y)

The expression:

x.sum().item()

turns tensor data into a Python scalar and then uses it for Python control flow.

That can force a break in graph capture.

Conceptually:

compiled graph A
      โ†“
return to Python
      โ†“
run unsupported / data-dependent logic
      โ†“
compiled graph B

The program may still work.

That is important.

A graph break is not automatically a correctness failure.

It is a boundary in what the compiler could capture together.

Too many graph breaks can destroy the optimization opportunity you expected from compilation.


5. Use fullgraph=True to expose hidden breaks

Normal compilation can tolerate graph breaks.

For debugging, that tolerance can hide the problem.

Try:

compiled = torch.compile(
    model,
    fullgraph=True,
)

Now PyTorch is being asked to capture the whole compiled region as one graph.

If something causes a graph break, you get a failure instead of silently falling back to multiple graph regions.

That makes fullgraph=True useful as a diagnostic mode.

Think of it as asking:

Can this region really be captured as one graph?

You do not necessarily need fullgraph=True in production.

Its value here is visibility.


6. Turn on compiler logging

Guessing is the wrong tool for compiler problems.

PyTorch exposes targeted compiler logs through TORCH_LOGS.

For graph breaks:

TORCH_LOGS="graph_breaks" python train.py

For recompilation:

TORCH_LOGS="recompiles" python train.py

For guards:

TORCH_LOGS="guards" python train.py

For dynamic-shape behavior:

TORCH_LOGS="dynamic,recompiles" python train.py

You can also enable some logging programmatically:

import torch

torch._logging.set_logs(
    graph_breaks=True,
    recompiles=True,
)

The goal is not to drown in compiler output.

The goal is to answer one question at a time.

Where did capture break?
Why did this graph compile again?
Which assumption failed?
Did shape variability cause it?

7. Guards are the compiler’s assumptions

Suppose the compiler sees:

x = torch.randn(32, 128)

When it creates optimized code, it may specialize based on facts about the execution it observed.

Those facts can include things such as:

shape
dtype
device
Python values
module state
other properties needed for correctness

The compiled graph is therefore not simply:

run this code forever

It is closer to:

run this graph while its assumptions remain valid

Those assumptions are checked with guards.

Conceptually:

input arrives
    โ†“
do guards still hold?
    โ†“
YES โ†’ reuse compiled graph
NO  โ†’ another graph may be needed

That is the bridge between guards and recompilation.


8. Recompilation is often a failed guard

Consider:

@torch.compile
def fn(x):
    return torch.sin(x) * 2


fn(torch.randn(32, 128))
fn(torch.randn(32, 256))
fn(torch.randn(32, 512))

If the compiled graph specialized on a dimension that changed, PyTorch may need another compiled version.

Run with:

TORCH_LOGS="recompiles" python example.py

Now the useful information is not merely:

recompiled

It is:

which guard failed?

A shape mismatch has a very different remedy from a changing Python constant or unsupported control flow.


9. Compilation has two kinds of cost

People often benchmark compiled code like this:

start = time.perf_counter()
y = compiled(x)
print(time.perf_counter() - start)

and conclude:

compilation is slower.

But a just-in-time compiler has at least two phases that matter:

cold path
    tracing + compilation + execution

warm path
    reuse compiled result + execution

Measure them separately.

import time

start = time.perf_counter()
compiled(x)
first_call = time.perf_counter() - start

start = time.perf_counter()
for _ in range(100):
    compiled(x)
steady = (time.perf_counter() - start) / 100

print({
    "first_call_s": first_call,
    "steady_state_s": steady,
})

On CUDA, remember the synchronization rules from the performance chapter.

A compiled model can have excellent steady-state speed and still be a poor choice for a short-lived process if compilation dominates total runtime.


10. Recompilation can turn cold cost into recurring cost

One expensive compilation may be acceptable.

Repeated compilation may not be.

Imagine:

shape A โ†’ compile
shape A โ†’ reuse
shape A โ†’ reuse
shape B โ†’ compile
shape C โ†’ compile
shape D โ†’ compile
shape E โ†’ compile

Now the system spends much more time in the cold path than you expected.

This is why TORCH_LOGS="recompiles" is so valuable.

The program may appear correct.

GPU kernels may even look fast.

But the wall-clock job can still be dominated by repeated compilation.


11. Dynamic shapes are a response to shape variation

PyTorch can compile with more dynamic shape behavior:

compiled = torch.compile(
    model,
    dynamic=True,
)

The idea is straightforward.

Instead of specializing aggressively to every observed size, generate code capable of handling a wider shape range where possible.

This can reduce recompilation for workloads such as:

variable sequence lengths
variable image sizes
variable batch sizes
ragged-ish application workloads

But dynamic shapes are not free magic.

They change the optimization problem.

So the correct workflow is:

prove recompilation is caused by shape variation
        โ†“
try dynamic-shape handling
        โ†“
measure compile count
        โ†“
measure steady-state performance
        โ†“
keep only if the real workload improves

Do not enable a compiler option merely because its name sounds relevant.


12. Static and dynamic are workload decisions

Suppose production always runs:

batch = 32
sequence = 2048

A highly specialized graph may be ideal.

Now suppose requests arrive as:

sequence 37
sequence 412
sequence 991
sequence 2048
sequence 83

A different trade-off appears.

The compiler question is not:

Are dynamic shapes better?

It is:

What distribution of shapes does my actual program see?

That is the same performance principle we have used throughout the book:

measure the workload you have, not the workload you imagine.


13. Narrow the compiled region

You do not need to compile an entire application.

Often the cleanest compiled boundary is the numerical hot path.

For example:

@torch.compile
def compiled_step(model, x):
    return model(x)

while leaving this outside:

logging
file I/O
metrics formatting
checkpoint writes
Python orchestration
progress bars
complex debugging code

A narrow boundary is easier to reason about.

It also reduces the amount of arbitrary Python the compiler must understand.


14. Debugging instrumentation can create compiler problems

Suppose you insert:

print(x.shape)
print(x.mean().item())

into a compiled function.

Those lines may change capture behavior.

That creates a nasty debugging loop:

program behaves strangely
    โ†“
add debugging prints
    โ†“
prints change graph capture
    โ†“
program behaves differently

One solution is to separate modes:

if debug:
    output = model(x)          # eager, instrumented
else:
    output = compiled_model(x) # compiled hot path

Another is to keep diagnostics outside the region being compiled.

Compiler debugging is one of the places where instrumentation boundaries matter.


15. torch.compiler.disable is an escape hatch

Sometimes one function simply does not belong in the compiled region.

PyTorch provides an explicit way to keep code eager:

import torch


@torch.compiler.disable
def noisy_python_helper(x):
    print("shape:", tuple(x.shape))
    return x

Then compiled code can call around or through eager regions as appropriate.

Do not treat this as failure.

A compiler boundary is an engineering choice.

The correct target is not:

100% of Python must compile

The target is:

important numerical regions compile well enough to improve the workload

16. Isolate the compiler layer

One of the most useful debugging techniques is to change the backend.

Start with:

compiled = torch.compile(
    model,
    backend="eager",
)

This still exercises Dynamo capture while avoiding the normal optimizing backend.

Then try:

compiled = torch.compile(
    model,
    backend="aot_eager",
)

Then the normal compiler path:

compiled = torch.compile(model)

The interpretation is roughly:

fails with backend="eager"
    โ†’ investigate Dynamo / graph capture / program interaction

works with eager, fails with aot_eager
    โ†’ investigate AOTAutograd path

works with aot_eager, fails with default backend
    โ†’ investigate backend lowering / Inductor path

This is not a perfect diagnostic theorem.

But it is an excellent way to reduce the search space.


17. Build a compiler probe

For small experiments, create a repeatable helper.

import time
import torch


def compile_probe(model, x):
    model = model.eval()

    with torch.inference_mode():
        eager = model(x)

    rows = []

    for backend in ["eager", "aot_eager", "inductor"]:
        candidate = torch.compile(model, backend=backend)

        start = time.perf_counter()
        with torch.inference_mode():
            out = candidate(x)
        first = time.perf_counter() - start

        torch.testing.assert_close(
            eager,
            out,
            rtol=1e-4,
            atol=1e-5,
        )

        rows.append({
            "backend": backend,
            "first_call_s": first,
        })

    return rows

This is deliberately simple.

It gives you a structured first question:

which compiler layer first changes correctness or fails?

For CUDA timing, use the synchronized benchmark helpers from the previous chapter.


18. Use tlparse when the model is too large for raw logs

TORCH_LOGS is excellent for a small reproducer.

For a large model, the output can become overwhelming.

PyTorch’s compiler tooling also supports trace-based inspection with TORCH_TRACE and tlparse.

The workflow is conceptually:

run program with compiler tracing
        โ†“
collect trace data
        โ†“
inspect high-level compilation frames
        โ†“
find graph breaks / recompiles / expensive frames
        โ†“
zoom into the relevant region

That is often better than reading thousands of raw log lines from the beginning.

The same debugging principle appears again:

Start with structure, then zoom into detail.


19. A graph break is not equally bad everywhere

Imagine two breaks.

Break A occurs once during setup before the hot numerical region.

Break B occurs repeatedly in the middle of every model step.

Those are not equivalent.

So the debugging target should not be:

zero graph breaks at any cost

Instead ask:

Where is the break?
How often is that path executed?
Does it split an important optimization region?
What is the measured performance cost?

Compiler diagnostics still need performance context.


20. Do not confuse compilation success with optimization success

This program can compile:

compiled = torch.compile(model)
compiled(x)

and still be a bad optimization.

Possible outcomes include:

compile succeeded, steady state faster
compile succeeded, steady state unchanged
compile succeeded, steady state slower
compile succeeded, memory increased
compile succeeded, many recompiles occur
compile succeeded, cold-start cost dominates

Therefore every compiler experiment needs at least two questions:

Did it compile correctly?
Did it improve the workload?

Those are separate results.


21. Keep a compiler evidence table

A tiny table is much better than memory.

configuration     first call    steady ms    recompiles    peak MiB
-------------------------------------------------------------------
eager             4 ms          4.0          0             1200
compile default    2.4 s         2.8          0             1260
compile dynamic    3.1 s         3.0          0             1280
compile shape-var  2.5 s         2.8          9             1260

The exact numbers do not matter.

The columns do.

Without them, statements such as:

torch.compile is faster

or:

dynamic shapes fixed it

are too vague to be useful.


22. Separate correctness, compilation and performance

This gives us a three-layer test:

LAYER 1 โ€” CORRECTNESS
Does eager execution produce the expected result?

LAYER 2 โ€” COMPILATION
Can the relevant region be captured and reused without pathological breaks/recompiles?

LAYER 3 โ€” PERFORMANCE
Does the compiled version improve the metric we care about?

Do them in that order.

Trying to optimize a model whose correctness is uncertain is a bad debugging strategy.

Trying to tune performance before you know recompilation behavior is often just as bad.


23. A practical compiler-debugging ladder

When compiled PyTorch behaves strangely, use this sequence:

    flowchart TD
    A[Run eager] --> B{Correct?}
    B -- No --> C[Fix ordinary model first]
    B -- Yes --> D[Compile and compare outputs]
    D --> E{Compiled correct?}
    E -- No --> F[Test backend=eager]
    F --> G[Test backend=aot_eager]
    G --> H[Test default/Inductor]
    E -- Yes --> I[Measure cold + warm performance]
    I --> J[Enable graph_breaks logs]
    J --> K[Enable recompiles logs]
    K --> L[Inspect guards / shape variation]
    L --> M[Test dynamic shapes if justified]
    M --> N[Re-measure real workload]
  

In text:

1. prove eager correctness
2. compare eager and compiled outputs
3. measure first-call and steady-state time separately
4. inspect graph breaks
5. inspect recompilations
6. identify failed guards
7. test dynamic shapes only if shape variation is real
8. isolate Dynamo / AOTAutograd / Inductor when failures remain
9. narrow the compiled region if useful
10. benchmark the real workload again

This is compiler debugging as engineering rather than folklore.


24. Challenge: create a graph break

Start with:

@torch.compile
def fn(x):
    return torch.sin(x) * 2

Verify that it runs.

Then introduce Python-visible data dependence:

@torch.compile
def fn(x):
    if x.sum().item() > 0:
        return torch.sin(x) * 2
    return torch.cos(x)

Run with:

TORCH_LOGS="graph_breaks" python example.py

Your goal is not merely to make the code compile again.

Your goal is to identify the exact boundary where graph capture changed.


25. Challenge: force recompilation

Try:

@torch.compile
def fn(x):
    return x.sin() + x.cos()


for n in [32, 64, 128, 256, 512]:
    fn(torch.randn(n, 128))

Run with:

TORCH_LOGS="recompiles" python example.py

Record:

number of calls
number of recompiles
first-call latency
latency for previously seen shapes
latency for new shapes

Then repeat with:

@torch.compile(dynamic=True)
def fn(x):
    return x.sin() + x.cos()

Do not assume the second version is better.

Measure it.


26. Challenge: identify the failing layer

Take a model that works eagerly but fails under normal torch.compile.

Try in order:

torch.compile(model, backend="eager")
torch.compile(model, backend="aot_eager")
torch.compile(model, backend="inductor")

Write down the first backend that fails.

That answer is already a much better bug report than:

torch.compile doesn’t work.


27. What this chapter adds to the book

Earlier chapters taught us to inspect:

tensor shapes
gradients
registered modules
data pipelines
activations
GPU timelines
memory

Compiler execution adds another hidden structure:

captured graphs
        โ†“
guards
        โ†“
graph reuse or recompilation

That structure is why compiled execution can behave differently from eager execution even when the source code barely changed.

But the debugging philosophy is exactly the same as before:

Make the hidden structure visible.

Once you can see graph breaks, guards and recompilations, torch.compile stops being a switch you either trust or distrust.

It becomes another system you can investigate.


Where we go next

There is one more failure mode before the final language-model build.

Sometimes nothing crashes.

The graph compiles.

The model trains.

All the unit tests pass.

And yet this week’s version is worse than last week’s.

The next chapter is about that problem:

training regressions and reproducible experiments

Because some bugs are not failures inside one run.

They are differences between runs.