torch.compile: what it wins and where to put it

torch.compile trades a one-time graph-compilation cost for a faster steady state. That trade is worth it whenever a model runs for many iterations — but the size of the win, and even whether it appears at all, depends on the architecture being compiled and on exactly which module torch.compile wraps.

What is torch.compile, and why use it? By default, PyTorch runs a model in “eager mode”: every operation (every matrix multiply, convolution, addition, …) is dispatched to the GPU one at a time, exactly as your Python code calls it. This is flexible and easy to debug, but it leaves performance on the table — each operation is a separate GPU kernel launch, with overhead, and intermediate results are written to and read back from GPU memory between every step even when several operations could be fused into one. torch.compile (introduced in PyTorch 2.0) looks at the model once, builds a computation graph out of it, and generates optimized GPU kernels for that graph — fusing operations together and cutting out unnecessary memory traffic. Building that optimized version takes extra time the first time it runs (the “1st iteration” or “compile” cost below), but every call after that reuses the compiled kernels and runs faster. Whether this trade is worth it, and how much it wins, is exactly what we measure here.

This study measures that trade in two settings: a standalone denoiser swept across architectures, GPUs, and problem sizes, and a full distributed plug-and-play (PnP) reconstruction where torch.compile can be applied at different points in the pipeline.

We look at three questions:

1. How much steady-state speedup does compiling a denoiser actually give, and does it depend on the architecture or the GPU?
2. Why does the speedup vary so much between architectures?
3. For a distributed PnP solver, does where you apply torch.compile matter?

Experimental setting

The denoiser sweep runs three architectures (DnCNN, DRUNet, UNet) eager and compiled, across 2D shapes (512/1024/2048) and 3D shapes (128³, 8×512×512, 16×512×512), on A100 and H100 GPUs. What we measure is deliberately narrow: one forward pass of the denoiser — no physics, no reconstruction loop — timed at each image size, so the compile effect is isolated to the network itself rather than mixed with anything else. The compiled arm is compile=“pre” with fullgraph=True and mode=“max-autotune”, on a single non-distributed GPU — with no distributed wrapper present, pre here just means “compiled”. Steady-state timing averages the last 5 timed denoiser calls per configuration.

The PnP sweep reconstructs a 4096² 2D tomography problem (100 projection angles) with DRUNet, on 1 and 2 GPUs, across four compile configurations: None (eager), pre, post, and fused — described in detail in the third section below. First-iteration timing includes the one-time compilation cost; stable-iteration timing is the mean of the last 5 timed iterations.

Two benchmark configurations produce everything below. They both live in benchmark_inference/configs/experiments/, and each figure on this page names the one it came from:

How much does compiling a denoiser speed things up?

Performance gains vary substantially across models and hardware. UNet sees the largest improvement, reaching up to 2.4× on the H100 and 1.6× on the A100 for 2D shapes. DRUNet sits at the opposite extreme, gaining almost nothing overall: the only configuration where compiling pays off is 2D shapes on the H100 (1.6×), while every other combination — A100 with 2D, A100 with 3D, and H100 with 3D — comes in under 1.04×. DnCNN falls between the two. This ranking — UNet gains most, DnCNN second, DRUNet least — holds across both GPUs; only the absolute magnitude of the speedup shifts.

This also shows a modern-architecture effect: for the same denoiser and the same shape, H100 consistently gets a bigger multiplicative speedup from compiling than A100 does (e.g. DnCNN at 512: 1.27x on A100 vs 2.03x on H100). Newer hardware doesn’t just run faster in eager mode — it has more headroom for the compiler’s kernel fusion to exploit.

This spread is not a hardware quirk. It comes from what each architecture asks the GPU to do.

Why does the gain vary so much?

Most points sit left of the ridge, meaning DnCNN and UNet are clearly memory-bound on this hardware. DRUNet’s 2D arithmetic intensity (≈178–181 FLOP/byte) sits just right of the ridge (≈147) — it is already close to compute-bound before compiling at all, which lines up with it seeing the smallest gain of the three (≈1.5–1.6x on H100).

Being memory-bound sets the ceiling on how much compilation can help — but it does not fully determine where each architecture lands under that ceiling. DnCNN has the lowest arithmetic intensity of the three (≈68 FLOP/byte, clearly the most memory-bound), yet it is UNet, not DnCNN, that sees the largest speedup (2.4x vs 2.1x on H100 2D). The 3D points make this especially clear: DnCNN and UNet land at the same arithmetic intensity there (≈205 FLOP/byte) but compile to very different speedups (1.14x for DnCNN vs 1.6x for UNet). So arithmetic intensity explains the broad left-to-right trend — and DRUNet’s poor result in particular — but the exact ranking between memory-bound architectures also depends on architecture-level details (how many small, fusable operations each network has, for instance), not on arithmetic intensity alone.

In short: torch.compile helps most where the model is furthest from saturating the GPU’s compute units, and least where it is already close to it.

For a distributed solver, does it matter where you compile?

The denoiser experiments above compile a standalone module. A full PnP reconstruction is different: the denoiser sits inside a distributed wrapper that tiles it across GPUs, and the compiler has to be told which part of that pipeline to target.

The PyTorch team’s own guidance on this is specific:

We recommend applying torch.compile to the highest-level function that doesn’t cause excessive problems. Typically, it is your train or eval step with the optimizer but without the loop, your top-level nn.Module, or some sub-nn.Modules. torch.compile specifically doesn’t handle distributed wrapper modules like DDP or FSDP very well, so consider applying torch.compile to the inner module passed to the wrapper.

Our PnP solver exposes exactly this choice through three compile modes:

• pre — compile(denoiser/physics), then distribute(...): compile first, then wrap the already-compiled module for distribution. This is the inner-module case the guidance above recommends.
• post — distribute(denoiser/physics), then compile(...): wrap for distribution first (this is where inter-GPU communication gets introduced), then compile the already-distributed wrapper. This is the case the guidance specifically warns against.
• fused — exactly like post (distribute first, compile after), just compiling more at once. post compiles the distributed denoiser and the distributed physics operator separately. fused instead compiles the whole per-iteration step in one region — compile(distribute(pnp_step)), gradient update, denoiser prox, and clamp all traced together.

At 1 GPU there is no distributed wrapper at all, so pre, post, and fused are not really three different things yet — with nothing to be “before” or “after”, all three compile the exact same non-distributed denoiser and physics operators. The small spread we see there (speedup 1.49–1.53x) is measurement noise, not a real effect of compile mode.

The difference shows up once distribution is active: at 2 GPUs, pre reaches 1.98x, ahead of both post (1.72x) and fused (1.76x). This matches what we see testing other configurations beyond what is charted here: the three modes sometimes land close together, but where they diverge, it is pre that comes out ahead. That is consistent with the PyTorch guidance — compiling the inner module before it is wrapped avoids handing the compiler a distributed wrapper it was not designed to trace well.

pre has a second, separate advantage visible in the same chart: its 1st-iteration (compile) cost is also lower once distribution is active — 14.98s at 2 GPUs, versus 26.78s for post and 24.89s for fused. This follows directly from what each mode actually compiles: pre only ever compiles the denoiser acting on a single tile/patch, before tiling and distribution exist. post and fused compile after the distributed wrapper is in place, so the graph they trace has to account for the whole distributed cycle, not one patch — a bigger graph takes longer to compile. So pre is not just faster at steady state here; it also pays a smaller one-time cost to get there.

Conclusions

The compile win is architecture-dependent, not just hardware-dependent. UNet gains the most from torch.compile (up to 2.44x on H100), DRUNet the least (≈1.0–1.6x), and DnCNN in between. This ranking holds on both A100 and H100 — only the absolute speedup shifts, and it consistently shifts upward on the newer H100.

Being memory-bound is necessary for a large gain, but not sufficient to predict the exact ranking. DRUNet’s high arithmetic intensity (already at or past the compute-bound ridge in 2D) correctly predicts its small gain. But between the two memory-bound architectures, UNet beats DnCNN despite not having the lower arithmetic intensity — in 3D they even share the same arithmetic intensity and still compile to different speedups. Architecture details beyond the roofline model also matter.

For distributed PnP, compiling the inner module before the distributed wrapper (compile=“pre”) is the safer default in what we’ve measured so far. It matches PyTorch’s own guidance on avoiding compilation of DDP/FSDP-style wrappers, it is the mode that has held up best as GPU count increases, and it also compiles faster to begin with, since it only ever traces a single denoiser patch rather than the whole distributed cycle.