GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2

Contribution

PyTorch 2 compiles models just in time: TorchDynamo captures the Python program as an FX graph and TorchInductor optimises it. Some ordinary code patterns make capture stop part-way. Each graph break sends execution back to eager Python, which adds synchronisation between CPU and GPU and cuts the graph the compiler can optimise. GraphMend removes fixable breaks automatically, by rewriting the program’s source before PyTorch sees it.

Key idea

GraphMend analyses the program’s syntax tree and recognises three patterns that cause graph breaks: control flow that depends on tensor values, Python side effects such as print and logging calls, and validation checks guarded by conditions. It applies a transformation only where it can show, before running the program, that the rewrite keeps its meaning. The developer does not refactor anything by hand, and PyTorch captures larger, uninterrupted graphs.

Results

From the arXiv version of the paper:

Measure	Result
Hugging Face models surveyed	195
Models with graph breaks	27 (13.8%)
Graph breaks eliminated	107 of 147 (73%)
Models fully fixed	21
Cold-start speedup	up to 26×, 5× on average
Steady-state forward pass	up to 1.39× faster

Measured on NVIDIA RTX 3090, A40 and H100 GPUs. Cold-start latency is the first forward pass; steady-state latency is the mean of the next nine passes.

Conclusion

The paper concludes that source-level analysis and transformation work well alongside PyTorch’s JIT compilation, and improve both usability and performance.

@misc{kashmira2025graphmend,
  title         = {GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2},
  author        = {Kashmira, Savini and Dantanarayana, Jayanaka and Sathiyalogeswaran, Thamirawaran and Flautner, Krisztian and Tang, Lingjia and Mars, Jason},
  year          = {2025},
  eprint        = {2509.16248},
  archivePrefix = {arXiv},
  note          = {Accepted at CGO 2027}
}