Overcoming the PyTorch Recompilation Dilemma in Real-Time Inference Workloads

How to Reduce the Runtime Interference of Recompilation Spikes

Photo by Mark de Jong on Unsplash

Just-in-time (JIT) model compilation is one of the most important and impactful methods for accelerating the runtime performance of PyTorch models. Contrary to the default eager execution mode — in which each line of Python code is processed independently — torch.compile converts the Python model definition into an optimized machine-level computation graph. The first time you execute a compiled model, the compiler traces the Python code and converts it into an intermediate graph representation (IR) using TorchDynamo, creating what is known to as an FX Graph. The FX Graph is then passed to a compiler backend (such as TorchInductor), which performs kernel fusion, out-of-order execution, and other techniques to generate machine code that is highly optimized for the target hardware. While the impact of PyTorch model compilation on runtime performance depends on the details of the model architecture and runtime environment (e.g., see the torch.compile benchmarking page), it is an essential tool for any modern-day AI/ML developer.

The Curse of Recompilations

One major drawback of torch.compile is the potential for costly recompilations throughout the lifecycle of the PyTorch application. While converting a PyTorch model into an intermediate graph representation, TorchDynamo must make multiple assumptions on the details of the underlying tensors, including shape, dtype, memory layout, constant values, and more. These assumptions are collected into a list of Dynamo guards — boolean checks that PyTorch creates to ensure that a previously compiled execution graph is safe to reuse for new inputs. When the dynamo guards are violated, a “recompilation” is triggered, generating a new compilation that accounts for the new input properties.

The frequency of recompilations depends heavily on the definition of the model, the details of the input data, the compilation configuration settings, and the runtime environment. For example, a model with fixed-shape inputs has a significantly lower risk of recompilation than a model processing dynamic tensor shapes. Note that when compiling a model designed for dynamic inputs, setting dynamic=True can help reduce the recompilation frequency.

One of the frustrating aspects of recompilations is that they can be extremely difficult to anticipate — let alone prevent. It is common for multiple recompilation events to occur during the first few dozen model passes — as the compilation calibrates to the specifics of the input data. However, there is no guarantee that recompilations will not occur much later in the application run. Often, recompilations will be triggered by rare and unique data samples that can show up sporadically during the application lifetime.

Another frustration developers often face is that recompilation behavior can vary greatly based on the runtime environment and PyTorch version. You may find yourself working awfully hard to circumvent recompilations in your current PyTorch environment, only to discover that your solution breaks when upgrading to a newer PyTorch version.

The PyTorch documentation acknowledges the inconvenience of recompilations and has a page dedicated to “dealing” with them. The page includes some tips for managing the occurrence of recompilations. But it does not offer a foolproof way to avoid them.

The Recompilation Dilemma

The potential for a costly recompilation mid-execution, poses a serious dilemma for PyTorch model developers: As compelling as the performance gains of model compilation may be, can we accept the possibility of a potentially lengthy disruption to our application? For offline applications, such as model-training workloads, an occasional spike in step time may be considered a small price to pay for the overall increase in training speed and reduction in costs. But in real-time inference settings with mission-critical response time requirements, such spikes could prove catastrophic. Consequently, it is not uncommon for development teams to stick with the slower, but more predictable eager-execution mode in such settings — leaving the opportunity for the runtime optimization associated with model compilation on the table.

In this post we address the difficult dilemma of recompilations on the inference pipeline. The question we ask is: Can we apply PyTorch compilation without exposing our inference application to a potentially catastrophic disruption. We will briefly explore several mitigations and alternatives, including:

• Preloading an existing compilation cache
• A hybrid eager-compiled execution approach
• Controlling compilation behavior via Dynamo configuration settings
• Seeding the graph compiler during a warm-up phase
• Executing recompilations on a designated background thread with minimal disruption to the inference function, and
• Fixing the computation graph using ahead-of-time (AOT) compilation

To facilitate our discussion, we will use a toy model and dataset specifically designed to trigger multiple recompilation events. We will run our model on an Amazon EC2 g6e.2xlarge instance with an NVIDIA L40S GPU running an AWS Deep Learning (Ubuntu 24.04) AMI with PyTorch (2.13). We will measure the effect of the recompilation events on inference performance and assess the impact of each mitigations. Importantly, the code we will share is intended for demonstrative purposes; Please do not rely on its correctness or optimality. The impact of the methods we will discuss will depend greatly on the details of the model, data, and runtime environment at hand. Please be sure to run your own evaluation before adopting them. Please do not interpret our mention of any library, platform, or API as an endorsement of its use.

A Toy GNN Model

For our demonstration, we build a simple graph neural network (GNN). GNNs are notorious for their highly dynamic shapes, making them highly prone to recompilations. The model below consists of a single node type and two relation types: “local” edges and “cross” edges. In the context of this post, familiarity with GNN semantics is not essential. The interested reader in referred to A Gentle Introduction to Graph Neural Networks for a brief overview on GNNs.

import torch
import torch.nn as nn

FEATURE_DIM = 16
HIDDEN_DIM = 768
MESSAGE_PASSING_LAYERS = 16

class MessagePassingModel(nn.Module):

    def __init__(self) -> None:
        super().__init__()
        # Encode each node feature vector as a 768-value hidden state.
        self.input_projection = nn.Linear(FEATURE_DIM, HIDDEN_DIM)

        # Each of the layers has one shared message transform and one node
        # update transform. Local and cross relations use the same learned
        # message transform, but keep separate edges and per-edge scales.
        self.message_projections = nn.ModuleList(
            nn.Linear(HIDDEN_DIM, HIDDEN_DIM)
            for _ in range(MESSAGE_PASSING_LAYERS)
        )
        self.update_projections = nn.ModuleList(
            nn.Linear(HIDDEN_DIM, HIDDEN_DIM)
            for _ in range(MESSAGE_PASSING_LAYERS)
        )

        # Project pooled 16-value context rows into the hidden-state space,
        # then map the combined graph representation to four output values.
        self.context_projection = nn.Linear(FEATURE_DIM, HIDDEN_DIM)
        self.output_projection = nn.Linear(HIDDEN_DIM, 4)

        # Eager pass counter
        self.eager_passes = 0

    def reset_eager_passes(self) -> None:
        self.eager_passes = 0

    def forward(
        self,
        x: torch.Tensor,
        local_edge_index: torch.Tensor,
        cross_edge_index: torch.Tensor,
        local_edge_scale: torch.Tensor,
        cross_edge_scale: torch.Tensor,
        context: torch.Tensor,
    ) -> torch.Tensor:
        # count eager passes
        if not torch.compiler.is_compiling():
            self.eager_passes += 1
        # One initial hidden state per node: [num_nodes, 768].
        state = self.input_projection(x).relu()
        for message_projection, update_projection in zip(
            self.message_projections, self.update_projections
        ):
            # edge_index[0] gathers one source-node state per edge. Each
            # gathered state is transformed and then weighted per edge.
            local_messages = (
                message_projection(state[local_edge_index[0]])
                * local_edge_scale
            )
            cross_messages = (
                message_projection(state[cross_edge_index[0]])
                * cross_edge_scale
            )

            # edge_index[1] identifies destinations. index_add sums
            # incoming messages at every destination node. The relations
            # aggregate independently, so either relation may be empty.
            local_sum = torch.zeros_like(state).index_add(
                0, local_edge_index[1], local_messages
            )
            cross_sum = torch.zeros_like(state).index_add(
                0, cross_edge_index[1], cross_messages
            )

            # A learned self/update term plus both message sums becomes the
            # next node state. Repeating this 16 times propagates
            # information across multiple graph hops.
            state = (
                update_projection(state)
                + local_sum
                + cross_sum
            ).relu()

        # Sum node and context rows into two graph-level 768-value vectors,
        # combine them, and predict four output values.
        return self.output_projection(
            state.sum(0) + self.context_projection(context).sum(0)
        )

Next, we define a dataset with 300 samples into which we purposely inject graphs of varying topologies.

import random

MIN_SIZE = 1_000
MAX_SIZE = 4_000
INJECT_ID = (0, 4, 12, 24, 40, 60, 84, 112, 144, 180)


def _empty_edges(device: torch.device) -> torch.Tensor:
    return torch.empty_strided(
        (2, 0), (0, 1), dtype=torch.long, device=device
    )


def _make_edges(
    nodes: int,
    count: int,
    generator: torch.Generator,
    device: torch.device
) -> torch.Tensor:
    if count == 0:
        return _empty_edges(device)
    return torch.stack((
        torch.randint(nodes, (count,), generator=generator, device=device),
        torch.randint(nodes, (count,), generator=generator, device=device),
    ))


def _apply_pattern(
    catalog_number: int | None,
    nodes: int,
    local_edges: int,
    cross_edges: int,
    context_rows: int,
) -> tuple[int, int, int, int]:
    # Non-catalog request: keep four independent, positive random counts.
    if catalog_number is None:
        return nodes, local_edges, cross_edges, context_rows
    if catalog_number == 0:
        # Both relations have the same edge count.
        return nodes, local_edges, local_edges, context_rows
    if catalog_number == 1:
        # Local relation empty
        return nodes, 0, cross_edges, context_rows
    if catalog_number == 2:
        # Cross relation empty
        return nodes, local_edges, 0, context_rows
    if catalog_number == 3:
        # Both relations empty
        return nodes, 0, 0, context_rows
    if catalog_number == 4:
        # Context empty
        return nodes, local_edges, cross_edges, 0
    if catalog_number == 5:
        # Local relation and context empty
        return nodes, 0, cross_edges, 0
    if catalog_number == 6:
        # Cross relation and context empty
        return nodes, local_edges, 0, 0
    if catalog_number == 7:
        # Nodes only: both relations and context are empty.
        return nodes, 0, 0, 0
    if catalog_number == 8:
        # Node-empty graph with context still present.
        return 0, 0, 0, context_rows
    if catalog_number == 9:
        # Fully empty graph and context.
        return 0, 0, 0, 0
    raise ValueError(f"unknown catalog pattern: {catalog_number}")


def _make_sample(
    device: torch.device,
    nodes: int,
    local_edges: int,
    cross_edges: int,
    context_rows: int,
    tensor_rng: torch.Generator,
) -> dict[str, torch.Tensor]:
    """Materialize a graph from its four row/edge counts."""
    return {
        "x": torch.randn(
            nodes, FEATURE_DIM, generator=tensor_rng, device=device
        ),
        "local_edge_index": _make_edges(
            nodes, local_edges, generator=tensor_rng, device=device
        ),
        "cross_edge_index": _make_edges(
            nodes, cross_edges, generator=tensor_rng, device=device
        ),
        "local_edge_scale": torch.randn(
            local_edges, 1, generator=tensor_rng, device=device
        ),
        "cross_edge_scale": torch.randn(
            cross_edges, 1, generator=tensor_rng, device=device
        ),
        "context": torch.randn(
            context_rows, FEATURE_DIM, generator=tensor_rng, device=device
        ),
    }


def make_dataset(
    device: torch.device, samples: int = 300, seed: int = 17
) -> list[dict[str, torch.Tensor]]:
    tensor_rng = torch.Generator(device=device).manual_seed(seed)
    shape_rng = random.Random(seed)
    catalog_by_request = {
        request_index: catalog_number
        for catalog_number, request_index in enumerate(INJECT_ID)
    }
    dataset: list[dict[str, torch.Tensor]] = []

    for request_index in range(samples):
        # Start with four independently selected positive counts.
        nodes = shape_rng.randint(MIN_SIZE, MAX_SIZE)
        local_edges = shape_rng.randint(MIN_SIZE, MAX_SIZE)
        cross_edges = shape_rng.randint(MIN_SIZE, MAX_SIZE)
        context_rows = shape_rng.randint(MIN_SIZE, MAX_SIZE)

        # Scheduled catalog requests override the random positive counts.
        nodes, local_edges, cross_edges, context_rows = _apply_pattern(
            catalog_by_request.get(request_index),
            nodes, local_edges, cross_edges, context_rows,
        )

        dataset.append(_make_sample(
            device,
            nodes=nodes,
            local_edges=local_edges,
            cross_edges=cross_edges,
            context_rows=context_rows,
            tensor_rng=tensor_rng,
        ))
    return dataset

In our baseline experiment, we calculate the average runtime of the GNN model on our dataset in standard eager mode:

import time

def infer_step(model: nn.Module, sample: dict[str, torch.Tensor]):
    torch.cuda.synchronize()
    start = time.perf_counter()
    model(**sample)
    torch.cuda.synchronize()
    return (time.perf_counter() - start) * 1_000

device = torch.device("cuda")
model = MessagePassingModel().to(device).eval()
dataset = make_dataset(device)
eager = [infer_step(model, sample) for sample in dataset]

# omit first sample from mean calculation
print(f"eager mean={sum(eager[1:])/(len(eager) - 1):.2f}ms")

The average step time of our baseline experiment running on our NVIDIA L40S GPU is 7.06 milliseconds.

The Model Compilation Tradeoff

In a previous post, we expanded on the benefits of applying JIT compilation using torch.compile. In the following code block, we apply it to our toy model and measure its impact on runtime. Given the dynamic nature of our model we configure dynamic=True. In anticipation of recompilation events, we configure the output torch logs to include “recompiles” and increase the compilation cache limit to 64 (default is 8). We will use the torch._dynamo.utils counters collection to keep track of recompilations:

import os
import shutil
import pathlib
import torch._dynamo as dynamo
import torch._dynamo.utils as dynamo_utils

def reset_compiler() -> None:
    # 1) Clear in-memory Dynamo compilation cache and counters.
    dynamo.reset()
    dynamo_utils.counters.clear()

    # 2) Use deterministic cache paths.
    cache_root = pathlib.Path("/tmp/.cache/compile")
    inductor_cache = cache_root / "inductor"
    triton_cache = cache_root / "triton"
    shutil.rmtree(cache_root, ignore_errors=True)
    inductor_cache.mkdir(parents=True, exist_ok=True)
    triton_cache.mkdir(parents=True, exist_ok=True)
    os.environ["TORCHINDUCTOR_CACHE_DIR"] = str(inductor_cache)
    os.environ["TRITON_CACHE_DIR"] = str(triton_cache)

    # 3) Clear generated Python module cache in this process.
    from torch._inductor.codecache import PyCodeCache
    PyCodeCache.cache_clear()


def unique_graph_count() -> int:
    """Number of distinct graphs Dynamo has compiled in this process."""
    return dynamo_utils.counters.get("stats", {}).get("unique_graphs", 0)

def infer_step_compiled(
        model: nn.Module, sample: dict[str, torch.Tensor]
    ):
    # Detect whether this inference caused Dynamo to compile a new graph.
    before_graphs = unique_graph_count()
    torch.cuda.synchronize()
    start = time.perf_counter()
    with torch.inference_mode():
        model(**sample)
    torch.cuda.synchronize()
    infer_time = (time.perf_counter() - start) * 1_000
    compiled_new_graph = unique_graph_count() > before_graphs
    return infer_time, compiled_new_graph

def run_experiment(model: torch.nn.Module, data: list):
    jit_times = []
    recompiles = []
    initial_compile_seen = False
    model.reset_eager_passes()

    for index, sample in enumerate(data):
        compiled_time, compiled_new_graph = infer_step_compiled(
            model, sample
        )
        jit_times.append(compiled_time)
        # The first new graph is initial compilation, not a recompile.
        # Every later new frame is an actual recompile.
        if compiled_new_graph and initial_compile_seen:
            recompiles.append(index)
        elif compiled_new_graph:
            initial_compile_seen = True

    excluded_indices = set(recompiles)
    stable_indices = [
        index
        for index in range(1, len(data))
        if index not in excluded_indices
    ]
    print(f"actual recompilations: {len(recompiles)}")
    print(f"actual recompile sample indices: {recompiles}")
    print(f"guard-miss eager fallbacks: {model.eager_passes}")
    if stable_indices:
        compiled_mean = (
            sum(jit_times[index] for index in stable_indices)
            / len(stable_indices)
        )
        print(
            f"compiled mean without recompilation: {compiled_mean:.2f}ms"
        )
        speedup = sum(eager[index] for index in stable_indices) / sum(
            jit_times[index] for index in stable_indices
        )
        print(f"speedup without recompilation: {speedup:.2f}x")
    else:
        print("speedup without recompilation: unavailable")
    if recompiles:
        average_compile_time = sum(
            jit_times[index] for index in recompiles
        ) / len(recompiles)
        max_compile_time = max(
            jit_times[index] for index in recompiles
        )
        print(f"average recompile time: {average_compile_time:.2f}ms")
        print(f"max recompile time: {max_compile_time:.2f}ms")
    else:
        print("average recompile time: unavailable")
        print("max recompile time: unavailable")

reset_compiler()
# log recompile events to the console
torch._logging.set_logs(recompiles=True)
# configure Dynamo compile cache limit
dynamo.config.cache_size_limit = 64
# compile model
compiled_model = torch.compile(model, dynamic=True, fullgraph=True)
run_experiment(compiled_model, dataset)

When running this script, the console fills up with multiple recompilation messages (abbreviated and reformatted for readability):

torch/_dynamo/guards.py:5396] 

[0/1] [__recompiles] Recompiling function forward in /home/ubuntu/test.py:38
triggered by the following guard failure(s):
cross_edge_index.size()[1] == local_edge_index.size()[1]  
# duck sizing added this equality because these variables had the same size
# (to avoid this specialization, set 
# torch.fx.experimental._config.use_duck_shape = False)

[0/2] [__recompiles] Recompiling function forward in /home/ubuntu/test.py:38
triggered by the following guard failure(s):
2 <= local_edge_index.size()[1] <= 2147483647  
# message_projection(state[local_edge_index[0]])  
# /home/ubuntu/test.py:58 in forward
# (the guard itself is not due user code but due to 0/1 specialization)
# to avoid specialization try torch._dynamo.decorators.mark_unbacked


actual recompilations: 10
actual recompile sample indices: [1, 4, 12, 24, 40, 60, 84, 112, 144, 180]
compiled mean without recompilation: 6.03ms
speedup without recompilation: 1.18x
average recompile time: 3837.45ms
max recompile time: 5568.60ms

Of the 300 input samples, ten triggered recompilation events. Their indices correspond to the unique graph topologies we injected into the dataset.

The tradeoff of model compilation is clearly evident: When isolating the warm compiled runs, we find that the performance increases by 18% compared to eager mode (6.03 ms vs. 7.06 ms). However the average recompilation time is over 500 times slower than the average baseline execution and the maximum recompilation time is nearly 800 times slower. Such severe tail-latency spikes are typically unacceptable in production inference environments.

Faced with this dilemma, many engineering teams will often abandon the torch.compile option and stick with eager execution. Coding assistants often recommend padding input tensors to fixed shapes or advise against JIT compilation for dynamic workloads.

If you are familiar with our previous posts, you will know that it is our belief that no stone shall be left unturned in our quest for performance optimization. In keeping with that conviction, the option of giving up on torch.compile as an optimization tool should be a last resort.

In the next sections we will explore a number of options for applying model compilation on dynamic graphs in production while mitigating its potentially destructive impact.

Option 1: Preload Compilation Cache

A standard method for reducing the overhead of model compilation is preloading the compilation cache from a previous runtime session using the torch.compiler.load_cache_artifacts API (see the torch.compile caching tutorial). By preloading a cache, a PyTorch workload can take advantage of compilation work that was performed by previous invocations of the application instead of compiling graphs from scratch. This can significantly reduce the overhead of recompilation in case of a new input topology for which a matching graph already exists in the cache.

However, it is important to note that the cache artifacts do not include the Dynamo runtime state or its Python guards. Consequently, this method may result in a substantial reduction of the compilation overhead, but it will not prevent the recompilations from being triggered. To simulate preloading the compilation cache, we rerun our script without resetting the compiler cache:

# disable compiler reset
# reset_compiler()
torch._logging.set_logs(recompiles=True)
dynamo.config.cache_size_limit = 64
# compile model
compiled_model = torch.compile(model, dynamic = True, fullgraph = True)
run_experiment(compiled_model, dataset)

The runtime results are captured below:

actual recompilations: 10
actual recompile sample indices: [1, 4, 12, 24, 40, 60, 84, 112, 144, 180]
compiled mean without recompilation: 6.03ms
speedup without recompilation: 1.18x
average recompile time: 553.84ms
max recompile time: 747.26ms

While recompilations still occur, the average compilation overhead drops by ~85% (from ~3,837 ms down to ~553 ms). Note that this optimization only benefits graph structures encountered in previous runs; entirely new input topologies will still trigger a full compilation cycle.

Option 2: Rely on Recompilation Cache Limit

PyTorch limits the number of recompilations to dynamo.config.cache_size_limit (which defaults to 8). Once the limit is reached, any new sample that violates a dynamo guard will fall back to eager execution. This allows us to — at the very least — cap the number of latency spikes we can expect. So long as this limit is reached within a fixed number of steps (which we might be able to control via seeding — see below), we are guaranteed not to hit any spikes for the remainder of the production run.

Importantly, the model will continue to execute the compiled variant on input samples that comply with the dynamo guards. If the majority of samples comply with the resultant dynamo guards, then the benefit of this hybrid eager-compiled scheme will surely be felt.

Get Chaim Rand’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

This method requires that we set fullgraph=False. In the following code block we set attempt this mitigation with the cache limit set to 1 and 2.

CACHE_SIZE_LIMITS = [1, 2]

for limit in CACHE_SIZE_LIMITS:
    print(f"\nRunning experiment with cache size limit {limit}")
    reset_compiler()
    torch._logging.set_logs(recompiles=True)
    dynamo.config.cache_size_limit = limit
    # fullgraph=False allows Dynamo to split and specialize, so cache limit
    # effects are visible as graph variants/recompilations.
    compiled_model = torch.compile(model, dynamic=True, fullgraph=False)
    run_experiment(compiled_model, dataset)

In the case of our toy experiment, the impact of this method is greatly dependent on the cache limit setting. In our dataset, the second sample is representative of the vast majority of samples. Thus, if we enable 2 compilations, we fall back to eager execution on only just 9 samples. Whereas, if we compile just the first sample, we fall back to eager execution on all of the other samples:

Running experiment with cache size limit 1
actual recompilations: 0
actual recompile sample indices: []
guard-miss eager fallbacks: 299
mean without recompilation: 6.82ms
speedup: 1.03x

Running experiment with cache size limit 2
actual recompilations: 1
actual recompile sample indices: [1]
guard-miss eager fallbacks: 9
mean without recompilation: 5.99ms
speedup: 1.18x

With a cache size limit of 2, we are able to match the same 18% speedup we saw in our baseline compilation experiment — with just a single recompilation spike.

Option 3: Customize Dynamo Configurations

Sometimes, the recompilation logs will provide some hints for how to program Dynamo so as to relax the guards that triggered them. Other times the guard messages might imply that a different configuration of the input tensors (e.g., a different memory layout) could prevent a recompilation. Unfortunately, as compelling as the notion of controlling the Dynamo guards may sound, as of the time of this writing, there is no foolproof configuration that is guaranteed to completely prevent recompilations. Moreover, the impact of a specific Dynamo configuration could vary based on the runtime environment.

In the case of our toy experiment the Dynamo recompilation includes the following hints:

# first recompile:
torch/_dynamo/guards.py:5142] [0/3] [__recompiles]     - 0/0:
cross_edge_index.size()[1] == local_edge_index.size()[1]  
# duck sizing added this equality because these variables had the same size
# (to avoid this specialization,
# set torch.fx.experimental._config.use_duck_shape = False)

# second recompile:
torch/_dynamo/guards.py:5142] [0/2] [__recompiles]     - 0/1:
2 <= local_edge_index.size()[1] <= 2147483647
# (user code shown is first use of this value--the guard itself is not due
# user code but due to 0/1 specialization in the framework; to avoid
# specialization try torch._dynamo.decorators.mark_unbacked(tensor, dim))

In the following code block we apply the recommended configurations:

import torch.fx.experimental._config as fx_config
from torch._dynamo.decorators import mark_unbacked

def mark_sample_unbacked(sample: dict[str, torch.Tensor]) -> None:
    """Mark each variable request dimension unbacked."""
    variable_dimensions = (
        (sample["x"], 0),
        (sample["local_edge_index"], 1),
        (sample["cross_edge_index"], 1),
        (sample["local_edge_scale"], 0),
        (sample["cross_edge_scale"], 0),
        (sample["context"], 0),
    )
    for tensor, dimension in variable_dimensions:
        mark_unbacked(tensor, dimension)

print("Running experiment with Dynamo configurations")
reset_compiler()
torch._logging.set_logs(recompiles=True)
fx_config.use_duck_shape = False
dynamo.config.cache_size_limit = 64
mark_sample_unbacked(dataset[0])
compiled_model = torch.compile(model, dynamic=True, fullgraph=True)
run_experiment(compiled_model, dataset)

In the case of our toy example, programming the Dynamo settings solves all of the recompilations:

actual recompilations: 0
actual recompile sample indices: []
guard-miss eager fallbacks: 0
compiled mean without recompilation: 5.96ms
speedup without recompilation: 1.18x

However, it is important to note that this method can be quite brittle and can largely depend on the ordering of input samples. We can easily demonstrate this by simply switching between the first and fourth samples in our catalogue:

def _apply_pattern(
    catalog_number: int | None,
    nodes: int,
    local_edges: int,
    cross_edges: int,
    context_rows: int,
) -> tuple[int, int, int, int]:
    # Non-catalog request: keep four independent, positive random counts.
    if catalog_number is None:
        return nodes, local_edges, cross_edges, context_rows
    if catalog_number == 3:
        # Both relations have the same edge count.
        return nodes, local_edges, local_edges, context_rows
    if catalog_number == 1:
        # Local relation empty
        return nodes, 0, cross_edges, context_rows
    if catalog_number == 2:
        # Cross relation empty
        return nodes, local_edges, 0, context_rows
    if catalog_number == 0:
        # Both relations empty
        return nodes, 0, 0, context_rows
    ...

In this case, not only to recompilation events remain, but new ones appear that are harder to fix (e.g., stride mismatch).

Recompiling function forward in /home/ubuntu/test.py:38
triggered by the following guard failure(s):
tensor 'cross_edge_index' stride mismatch at index 0. expected 0, actual 3888

actual recompilations: 6
actual recompile sample indices: [1, 4, 12, 40, 60, 84]
guard-miss eager fallbacks: 0
compiled mean without recompilation: 5.98ms
speedup without recompilation: 1.19x
average recompile time: 4729.22ms
max recompile time: 5667.14ms

Option 4: Seed the JIT Compiler

Another strategy is to seed the compiler with all possible sample topologies in a warmup phase before activating the inference pipeline. In the case of our toy experiment, we can create a function that iterates through all 10 graph topologies and runs the model on each one:

def warmup_dataset(device: torch.device) -> list[dict[str, torch.Tensor]]:
    """Create one deterministic representative for each catalog pattern."""
    tensor_rng = torch.Generator(device=device).manual_seed(1_017)
    shape_rng = random.Random(1_017)
    warmup: list[dict[str, torch.Tensor]] = []

    for catalog_number in range(len(INJECT_ID)):
        nodes = shape_rng.randint(MIN_SIZE, MAX_SIZE)
        local_edges = shape_rng.randint(MIN_SIZE, MAX_SIZE)
        cross_edges = shape_rng.randint(MIN_SIZE, MAX_SIZE)
        context_rows = shape_rng.randint(MIN_SIZE, MAX_SIZE)
        nodes, local_edges, cross_edges, context_rows = _apply_pattern(
            catalog_number, nodes, local_edges, cross_edges, context_rows
        )
        warmup.append(_make_sample(
            device,
            nodes=nodes,
            local_edges=local_edges,
            cross_edges=cross_edges,
            context_rows=context_rows,
            tensor_rng=tensor_rng,
        ))
    return warmup

print("Running experiment with warmup dataset")
warmup = warmup_dataset(device)
reset_compiler()
torch._logging.set_logs(recompiles=True)
dynamo.config.cache_size_limit = 64
# compile model
compiled_model = torch.compile(model, dynamic = True, fullgraph = True)
# warmup compiled model
with torch.inference_mode():
    for sample in warmup:
        compiled_model(**sample)
run_experiment(compiled_model, dataset)

Since all of the catalogue examples are seen during the warmup phase, we encounter no recompilation spikes during the inference loop:

actual recompilations: 0
actual recompile sample indices: []
guard-miss eager fallbacks: 0
compiled mean without recompilation: 5.94ms
speedup without recompilation: 1.19x

Of course, the success of this method assumes that we have full knowledge of all of the possible input topologies. It also assumes that we have the ability to program each possibility — which may prove difficult when a recompilation is triggered by an input sample’s specific memory layout. Once again, in some production settings, the likelihood of missing even one input type, could disqualify this option.

Option 5: Asynchronous Background Compilation

The trouble with recompiles stems from the fact that PyTorch’s JIT compiler is a blocking operation that can be extremely compute-intensive. If a Dynamo guard is triggered, the inference operation is suspended while a new compiled model is generated. In this section, we explore the possibility of programming the recompilations to be non-blocking so that we can use torch.compile without needing to worry about significant interference to the inference latency. The general scheme we propose is:

• Maintain two versions of the model: eager and compiled.
• Create a background thread designated for performing recompilations while inference executes on the foreground thread.
• So long as a Dynamo guard is not triggered, the foreground thread uses the compiled model.
• If a Dynamo guard is triggered, the foreground thread immediately reverts to the eager model and dispatches the violating sample to the background thread for recompilation.
• When background compilation finishes, the background thread signals the foreground thread, which seamlessly returns to using the compiled model.

The key to this approach is to minimize the interference of the background compilation on the eager inference of the foreground thread. Compilation can be a resource-intensive operation. The proposed solution assumes that there are sufficient resources (e.g., GPU memory) to run both processes concurrently.

Note, that we run both threads in a single process. This is required in order to ensure that the updates to the Dynamo runtime state during recompilation are applied to the foreground thread. Were we to implement the recompilation in a separate process, we would need to somehow carry the updated runtime state back to the foreground process.

However, Python uses a Global Interpreter Lock (GIL) to serialize execution across multi-threaded Python workloads. In our case, the compute heavy compilation worker can hold the GIL for extended durations, starving the main inference thread of CPU time and leading to high tail latency during eager fallback execution. (Note that as of the time of this writing, PyTorch does not fully support free-threaded, no-GIL Python.)

In the following code block, we attempt to reduce the interference from the background thread in a number of ways:

1. We pin the recompilation thread to a single CPU core
2. We reduce the thread’s scheduling priority
3. We restrict the number of compiler threads to 1
4. We program the GIL switch interval

import contextlib
import queue
import sys
import threading

class AsyncCompileWrapper:
    """A PyTorch model wrapper for asynchronous background recompiles."""

    def __init__(
        self, model: torch.nn.Module, gil_interval_ms: float | None = None
    ) -> None:
        self.eager = model.eval()
        self.compiled = torch.compile(model, fullgraph=True, dynamic=True)
        self.gil_interval_ms = gil_interval_ms
        self.jobs = queue.Queue(maxsize=1)
        # The state of this event determines whether foreground should try
        # compiled inference.
        self.compiling = threading.Event()
        # Count completed background compilations.
        self.completed = 0
        # Run background recompilation worker thread.
        self.worker = threading.Thread(target=self._worker, daemon=True)
        self.worker.start()

    @contextlib.contextmanager
    def _recompile_error(self, enabled: bool):
        # The foreground enables this so a guard miss becomes an exception.
        # The worker disables it to compile the miss.
        previous = dynamo.config.error_on_recompile
        dynamo.config.error_on_recompile = enabled
        try:
            yield
        finally:
            dynamo.config.error_on_recompile = previous

    def _worker(self) -> None:
        import torch._inductor.config as inductor_config
        # Pin only this Linux worker thread to one CPU
        allowed_cpus = os.sched_getaffinity(0)
        os.sched_setaffinity(0, {min(allowed_cpus)})

        # Lower the worker thread's scheduling priority.
        os.setpriority(os.PRIO_PROCESS, threading.get_native_id(), 10)

        # Restrict compiler CPU parallelism.
        inductor_config.compile_threads = 1

        while (sample := self.jobs.get()) is not None:
            previous = sys.getswitchinterval()
            if self.gil_interval_ms is not None:
                # Program GIL switch interval.
                sys.setswitchinterval(self.gil_interval_ms / 1_000)
            with torch.inference_mode(), self._recompile_error(False):
                self.compiled(**sample)
            self.completed += 1
            # Foreground may now try the compiled model again.
            self.compiling.clear()
            sys.setswitchinterval(previous)
            self.jobs.task_done()
        self.jobs.task_done()

    def __call__(self, sample: dict[str, torch.Tensor]):
        if self.compiling.is_set():
            # A worker already owns one compilation job,
            # so run this request eagerly while it completes.
            return self.eager(**sample), False
        try:
            # The following fails on a cache miss
            with self._recompile_error(True):
                output = self.compiled(**sample)
            return output, True
        except dynamo.exc.RecompileError:
            # Compile this triggering sample in background thread
            self.compiling.set()
            self.jobs.put(sample)
            return self.eager(**sample), False

    def close(self) -> None:
        self.jobs.put(None)
        self.jobs.join()
        self.worker.join(timeout=5)

def infer_step_async(
    model: AsyncCompileWrapper, sample: dict[str, torch.Tensor]
):
    torch.cuda.synchronize()
    start = time.perf_counter()
    with torch.inference_mode():
        _, compiled = model(sample)
    torch.cuda.synchronize()
    infer_time = (time.perf_counter() - start) * 1_000
    return infer_time, compiled


def run_async_experiment(
    model: AsyncCompileWrapper, data: list[dict[str, torch.Tensor]]
) -> None:
    eager_times = []
    compiled_times = []

    for index, sample in enumerate(data):
        async_time, compiled = infer_step_async(model, sample)
        # The first request includes the initial compilation/setup cost,
        # so it is not representative of either steady-state serving path.
        if index == 0:
            continue
        if compiled:
            compiled_times.append(async_time)
        else:
            eager_times.append(async_time)

    print(
        "eager runs: "
        f"{len(eager_times)}, compiled runs: {len(compiled_times)}"
    )
    eager_mean = (
        sum(eager_times) / len(eager_times)
        if eager_times
        else None
    )
    compiled_mean = (
        sum(compiled_times) / len(compiled_times)
        if compiled_times
        else None
    )
    eager_max = max(eager_times) if eager_times else None
    compiled_max = max(compiled_times) if compiled_times else None
    print(
        "execution time, excluding the first request: "
        f"eager={eager_mean:.2f}ms" if eager_mean is not None
        else "average execution time: eager=unavailable",
        f", compiled={compiled_mean:.2f}ms" if compiled_mean is not None
        else ", compiled=unavailable",
        f", eager max={eager_max:.2f}ms" if eager_max is not None
        else ", eager max=unavailable",
        f", compiled max={compiled_max:.2f}ms" if compiled_max is not None
        else ", compiled max=unavailable",
        sep="",
    )
    print(f"background compilations completed: {model.completed}")

dynamo.config.cache_size_limit = 64
for gil_interval_ms in [None, 0.1]:
    reset_compiler()
    async_base_model = MessagePassingModel().to(device).eval()
    async_model = AsyncCompileWrapper(async_base_model, gil_interval_ms)
    print(f"gil_interval_ms = {gil_interval_ms}")
    run_async_experiment(async_model, dataset)
    async_model.close()

In the case of our toy example, the most impactful setting turned out to be the GIL switch interval. Without it, the average eager execution time is 64.85 ms (~9 times slower than baseline) but with a peak latency of 1579.97 ms. When setting the GIL switch interval to 0.1, the average eager execution time drops to 35.30 ms (9 times slower than baseline) and the maximum latency to 317.34 ms. Whether these results are sufficient will depend on the product latency requirements.

Option 6: The Ahead-of-Time (AOT) Compilation Alternative

In addition to the torch.compile JIT compiler, PyTorch supports ahead-of-time (AOT) compilation via its torch.export API. In a recent post, we demonstrated the use of torch.export for model deployment. Naturally, AOT compilation does not suffer from the possibility of surprise recompiles. If an exported model encounters an input sample that violates its assumptions — it will fail. The key is configuring the export specifications with appropriate input bounds up front.

The output of torch.export is intended to be target-platform agnostic — one consequence of this being that it does not include the same platform specific optimizations that torch.compile does. However, we can apply the AOT version of TorchInductor to the output of torch.export to achieve the same type of platform level optimizations. Please see the torch._inductor.aoti_compile_and_package API documentation for details.

As discussed in our previous post, torch.export tends to be less flexible than torch.compile and its use may require some code refactoring. In the case of our toy model, no model changes were required:

from torch.export import Dim

def run_aot_experiment(
    base_model: nn.Module, data: list[dict[str, torch.Tensor]]
) -> None:
    """Export and run an AOT-compiled package on the dataset."""

    nodes_dim = Dim("nodes", min=0, max=6_000)
    local_edges_dim = Dim("local_edges", min=0, max=6_000)
    cross_edges_dim = Dim("cross_edges", min=0, max=6_000)
    context_dim = Dim("context", min=0, max=6_000)
    dynamic_shapes = (
        {0: nodes_dim},
        {1: local_edges_dim},
        {1: cross_edges_dim},
        {0: local_edges_dim},
        {0: cross_edges_dim},
        {0: context_dim},
    )

    exported = torch.export.export(
        base_model, tuple(data[0].values()), dynamic_shapes=dynamic_shapes
    )

    package_path = pathlib.Path("/tmp/.cache/compile/model_aot.pt2")
    if package_path.exists():
        package_path.unlink()
    built = torch._inductor.aoti_compile_and_package(
        exported, package_path=str(package_path)
    )
    serving_model = torch._inductor.aoti_load_package(built)

    def infer_step_aot(sample: dict[str, torch.Tensor]) -> float:
        torch.cuda.synchronize()
        start = time.perf_counter()
        serving_model(*sample.values())
        torch.cuda.synchronize()
        return (time.perf_counter() - start) * 1_000

    aot_times = [infer_step_aot(sample) for sample in data]
    stable_times = aot_times[1:]
    aot_mean = sum(stable_times) / len(stable_times)
    aot_max = max(stable_times)
    print(f"aot mean={aot_mean:.2f}ms")
    print(f"aot max={aot_max:.2f}ms")
    speedup = sum(eager[1:]) / sum(stable_times)
    print(f"aot speedup vs eager={speedup:.2f}x")


aot_base_model = MessagePassingModel().to(device).eval()
run_aot_experiment(aot_base_model, dataset)

Note, that programming this solution required specification of the boundaries of the dynamic dimensions.

The AOT-compiled implementation results in mean step time of 5.8 miliseconds which amounts to a 21% improvement over the eager execution baseline.

When the model is export-compatible and the input data shapes have clearly defined boundaries, AOT compilation could be the perfect alternative to the recompile-prone JIT option.

AOT with Eeager Fallback

If you want to adopt AOT compilation but are concerned about edge cases or outliers violating your input-contract assumptions, consider wrapping the compiled model with an eager fallback mechanism. If an input falls outside the expected bounds, the system can gracefully fall back to eager execution. Be sure to log detailed diagnostics whenever a fallback occurs so you can update your AOT shape definitions accordingly in future releases.

Summary

Unexpected recompilation spikes represent a potential nightmare scenario for latency-sensitive production environments, often driving engineering teams to abandon torch.compile in favor of the slower but predictable eager execution.

In this post we covered several methods for reducing or eliminating the recompilation overhead. There is no single silver bullet when it comes to PyTorch recompilations. Each approach has its own tradeoffs, advantages, and limitations. We hope that by tailoring techniques such as these to your specific model architecture, data distribution, and application requirements, you can capture the runtime speedups of PyTorch compilation without exposing your production workload to catastrophic latency spikes.