Inference Systems · Graph Compilation

Graph Breaks, Fallback Paths, and Numerical Drift

What happens at the boundary where tracing stops, why the eager fallback does not produce bit-identical results, and how a runtime keeps two execution paths consistent.

A graph compiler sees a program through a narrow aperture. It observes operations on tensors and nothing else. A Python loop whose trip count depends on a tensor’s contents, a call into a library that was not registered as a compiler-visible operator, a print of an intermediate value, a mutation of a module attribute, an exception handler that inspects a shape: all of these are invisible to the tracer, and the trace cannot continue past them. The compiler stops, emits what it has, and resumes tracing after the offending statement. The resulting program is a sequence of compiled regions separated by the interpreter.

Each boundary is a graph break. The name suggests a defect, and in one sense it is: a model compiled as one region can be fused, planned, and replayed as a unit, and a model compiled as forty regions can do none of those across the boundaries. But the boundary is not a bug to be fixed. It is the mechanism by which a compiler that cannot see everything still compiles what it can, and the interesting engineering is in what happens at the seam.

How tracing produces a boundary

Tracing runs the model with example inputs and records the operations performed. It is a partial evaluation: the tracer intercepts tensor operations, builds graph nodes for them, and computes their results so that subsequent Python control flow sees real values. When the Python code reaches a construct the tracer cannot record, it has a choice. It can fail the whole compilation, which makes the compiler useless on any model with a conditional. Or it can stop recording, let the interpreter execute the construct with the values it has, and start a new recording afterwards.

The second choice is what every practical system does, and it has a consequence that is easy to miss: the boundary is a synchronisation point. To let the interpreter see a tensor’s value, the tracer must read that value back from the device to the host, and that read is blocking. A graph with twenty breaks has twenty host round trips that no amount of fusion can remove, and each of them serialises the device pipeline. This is the dominant cost of a break, far exceeding the loss of fusion opportunities.

The reasons a break occurs fall into a few classes. Data-dependent control flow is the largest: an if on a tensor value, a while whose condition reads a tensor, or a loop over a shape that was computed from an input. Host-side effects are the second: logging, printing, asserting on a tensor’s contents, or writing to a file. Unsupported operations are the third: a call into a custom extension that has no registered decomposition. And mutation of Python state is the fourth: assigning to an attribute of a module that the compiled region closed over.

def forward(x):
    y = model.block1(x)              # traced into region A
    if y.sum() > 0:                  # data-dependent: break here
        y = model.block2(y)          # region B
    else:
        y = -y                       # region B, other branch
    return model.head(y)             # region C

Nothing in that function is unusual, and the break is unavoidable: the branch depends on a value, so the compiled region cannot contain both arms without a conditional node, and most tracers do not build one. The three regions are each compiled and fused independently, and the two host reads between them dominate the runtime of a small model.

What the boundary costs beyond the sync

The first structural cost is layout. A compiled region chooses physical orders for its values based on its own operators and its own consumers. It cannot see the next region, so its output layout is chosen from the region’s own preference, and the next region’s input layout is chosen from its own. Where the two disagree, the boundary needs a conversion. In an unbroken graph the layout assignment would have propagated the constraint across the seam and eliminated the copy; here it cannot.

The second cost is memory. Each region gets its own arena and plan. A tensor that lives across a boundary must be materialised into memory owned by neither region’s arena, or into an arena shared by both, and either way the planner loses the ability to reuse that buffer for anything else. The peak footprint of a forty-region model is routinely larger than the footprint of the same model as one region, because the planner can only reuse within a region.

The third cost is the launch structure. A region can be captured and replayed as a unit, but the sequence of regions cannot, because the interpreter between them issues its own work. Replay recovers the per-launch cost within a region and cannot touch the cost between regions.

The cumulative effect is that a model with many breaks behaves like an eager model with some local fusion, which is roughly what it is. The gain from compilation scales with the size of the average region, not with the size of the model.

Two execution paths and the drift between them

A runtime with a compiler almost always retains an uncompiled path. It is needed for inputs that fail the guards, for debugging, and for the many small operations where compilation overhead exceeds the gain. The existence of two paths creates an obligation: they must agree.

They do not agree bit for bit, and the reasons are all in the arithmetic. A compiled region may fuse a multiply and an add into a single fused-multiply-add, which rounds once instead of twice. It may accumulate a reduction in a different order, because it tiles the reduction and each tile’s partial sum is added in tile order rather than in index order. It may keep an intermediate in a register at a wider precision than the eager path’s memory buffer. It may vectorise a loop in a way that changes the grouping of a sum. None of these is a compiler bug; each is a legitimate consequence of the transformation, and each changes the result in the last few bits.

The magnitude is bounded by the arithmetic. A float32 value has 24 bits of significand, so a relative error of one unit in the last place is about $6 \times 10^{-8}$. A reduction over $n$ terms in a different order accumulates error proportional to $\sqrt{n}$ in the favourable case and to $n$ in the worst, so a reduction over a million terms can differ in the fourth or fifth significant digit. For a model whose output is a probability distribution, that is invisible. For a model whose output is compared against a threshold, or fed into a control loop, or used to decide a branch in a subsequent stage, it can flip a decision.

This is why the standard practice is to compare compiled against eager outputs with a tolerance rather than for equality, and why the tolerance has to be chosen with the model’s sensitivity in mind. A tolerance of $10^{-5}$ on a float32 model catches genuine errors, such as a mis-fused kernel or a wrong offset, and admits the legitimate drift.

Keeping the seam honest

The engineering that keeps a broken graph correct is mostly about the state that crosses the boundary. A value produced in one region and consumed in the next must be handed over with its identity intact: if the consumer is a view operation that expects to alias, the handover must preserve the aliasing or the consumer will silently read a copy. A tensor mutated in the interpreter between regions must be visible to the region that follows, which means the compiler must either treat the mutation as an ordering constraint or decline to compile across it.

The same reasoning applies to randomness and to any operator with hidden state. A random draw inside a compiled region must be seeded from a value the graph carries, or the compiled path and the eager path will diverge not by rounding but by content. A stateful operator such as a batch-norm in training mode must be excluded from the compiled region entirely, because its update is a side effect the graph cannot express.

The most reliable way to find these problems is a differential test: run the same inputs through both paths and compare, with a tolerance, on every distinct input shape the deployment sees. That test does not prove the paths agree, because it cannot cover the input space, but it catches the structural mistakes, which are the ones that produce large differences rather than last-bit ones. A difference that exceeds the tolerance by orders of magnitude is a boundary bug; a difference in the seventh digit is arithmetic.

The seam is where the model ends

The final observation is that a graph break is not only a compiler limitation. It marks the place where the model stops being a pure function of its inputs, and a compiler that could remove every break would still be compiling a program whose behaviour depends on host state it cannot see. Data-dependent control flow, host-side effects, and mutable state are properties of the program the model is embedded in, not defects of the graph format.

The realistic goal is therefore to make the regions large and the seams cheap, not to eliminate them. That means preferring tensor operations over scalar control flow where the model permits it, registering custom operators so the tracer can record them, and keeping host reads out of the hot path. Each of those increases the size of the average region, and the benefit of every transformation applied to it, from fusion through replay, is proportional to that size.

References

  1. PyTorch documentation — Troubleshooting torch.compile — PyTorch project
  2. PyTorch documentation — TorchDynamo Overview — PyTorch project
  3. JAX documentation — Just-in-time compilation — Google DeepMind
  4. IEEE 754-2019: Standard for Floating-Point Arithmetic — IEEE