about
GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2 (arxiv.org)
3 points by matt_d on Sep 25, 2025 | hide | past | pdf | discuss on HN

In plain words: GraphMend rewrites code to remove spots where PyTorch's compiler gives up and falls back to Python, only when it can prove behavior is unchanged. It removed 73% of the breaks across 27 models, clearing all in 21 and speeding startup up to 26 times.

Abstract

This paper presents GraphMend, a compiler technique that automatically fixes FX graph breaks in PyTorch 2 programs. Although PyTorch 2 introduced TorchDynamo and TorchInductor to enable just-in-time graph compilation, certain code patterns still cause graph breaks that force execution to fall back to Python eager mode, introducing costly CPU-GPU synchronization and reducing optimization opportunities. Our investigation of 195 Hugging Face models reveals that 13.8% of models exhibit graph breaks. GraphMend automatically eliminates fixable breaks through source-level program analysis and transformations. It analyzes AST-level program structure to identify graph-break patterns and applies transformations only when their semantic preservation can be statically established. These transformations enable PyTorch to capture larger, uninterrupted FX graphs without manual refactoring by developers. We evaluate GraphMend on all 27 models found to exhibit graph breaks in our investigation. GraphMend eliminates 107 of 147 graph breaks (73%), fully fixing all breaks in 21 models. In our experiments on NVIDIA GPUs, GraphMend achieves up to 26x cold-start speedup, 5x on average, and up to 1.39x steady-state forward pass speedup. These results demonstrate that semantics-aware source-level analysis and transformation are effective complements to PyTorch's dynamic JIT compilation pipeline, substantially improving both usability and performance.

Savini Kashmira, Jayanaka Dantanarayana, Thamirawaran Sathiyalogeswaran, Krisztian Flautner, Lingjia Tang, Jason Mars
arXiv:2509.16248 · cs.PL, cs.LG, cs.SE · submitted Sep 17, 2025 · updated Sep 3, 2026
abstract · pdf · html · Accepted for publication at the International Symposium on Code Generation and Optimization (CGO 2027)

add comment on HN