# GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2

arXiv: https://arxiv.org/abs/2509.16248 · HTML: https://mlsanity.com/paper/2509.16248/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 68 references that are in the map.
- First posted on arXiv 17 Sep 2025.

## Builds on (most cited first, 26 of 26)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Robust Speech Recognition via Large-Scale Weak Supervision](https://mlsanity.com/paper/2212.04356/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection](https://mlsanity.com/paper/2303.05499/) (2023)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [Chronos: Learning the Language of Time Series](https://mlsanity.com/paper/2403.07815/) (2024)
- [Large-scale Contrastive Language-Audio Pretraining with Feature Fusion and Keyword-to-Caption Augmentation](https://mlsanity.com/paper/2211.06687/) (2022)
- [Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models](https://mlsanity.com/paper/2311.07919/) (2023)
- [BioGPT: Generative Pre-trained Transformer for Biomedical Text Generation and Mining](https://mlsanity.com/paper/2210.10341/) (2022)
- [LayoutLMv3: Pre-training for Document AI with Unified Text and Image Masking](https://mlsanity.com/paper/2204.08387/) (2022)
- [Florence-2: Advancing a Unified Representation for a Variety of Vision Tasks](https://mlsanity.com/paper/2311.06242/) (2023)
- [Large-Scale Chemical Language Representations Capture Molecular Structure and Properties](https://mlsanity.com/paper/2106.09553/) (2021)
- [Jasper and Stella: distillation of SOTA embedding models](https://mlsanity.com/paper/2412.19048/) (2024)
- [Glow: Graph Lowering Compiler Techniques for Neural Networks](https://mlsanity.com/paper/1805.00907/) (2018)
- [MLIR: A Compiler Infrastructure for the End of Moore's Law](https://mlsanity.com/paper/2002.11054/) (2020)
- [Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python](https://mlsanity.com/paper/2112.08429/) (2021)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning](https://mlsanity.com/paper/1903.01855/) (2019)
- [HydraServe: Minimizing Cold Start Latency for Serverless LLM Serving in Public Clouds](https://mlsanity.com/paper/2502.15524/) (2025)
- [AutoGraph: Imperative-style Coding with Graph-based Performance](https://mlsanity.com/paper/1810.08061/) (2018)
- [The Jaseci Programming Paradigm and Runtime Stack: Building Scale-out Production Applications Easy and Fast](https://mlsanity.com/paper/2305.09864/) (2023)
- [MTP: A Meaning-Typed Language Abstraction for AI-Integrated Programming](https://mlsanity.com/paper/2405.08965/) (2024)
- [PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch](https://mlsanity.com/paper/2503.19779/) (2025)
- [Extending Data Spatial Semantics for Scale Agnostic Programming](https://mlsanity.com/paper/2504.03109/) (2025)

## Built on by (most cited first, 1 of 1)

- [Prompt Less, Smile More: MTP with Semantic Engineering in Lieu of Prompt Engineering](https://mlsanity.com/paper/2511.19427/) (2025)

## Abstract (arXiv)

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.