GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2

2025/09/17 by Savini Kashmira, Kashmira, Savini, Jayanaka Dantanarayana +12 · 1 voice · 1 citation
Computer Science · #Call graph #Code refactoring #Compiler #Control flow #Control flow graph #Data-flow analysis #Graph #Graph rewriting #Model-Driven Software Engineering Techniques #Natural Language Processing Techniques #Python (programming language) #Software Testing and Debugging Techniques #Source code #cs.LG #cs.PL #cs.SE

paper · pdf · doi:10.48550/arxiv.2509.16248

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2025/09/17 · arxiv published 2025/09/17 · openalex created_date 2025/10/10 · arxiv updated 2026/06/29 · openalex updated_date 2026/08/01

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.

Citations

• ServerlessLoRA: Minimizing Latency and Cost in Serverless Inference for LoRA-Based LLMs
• Extending Data Spatial Semantics for Scale Agnostic Programming
• PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch
• HydraServe: Minimizing Cold Start Latency for Serverless LLM Serving in Public Clouds
• Jasper and Stella: distillation of SOTA embedding models
• MTP: A Meaning-Typed Language Abstraction for AI-Integrated Programming
• Chronos: Learning the Language of Time Series
• Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models
• Florence-2: Advancing a Unified Representation for a Variety of Vision Tasks
• The Jaseci Programming Paradigm and Runtime Stack: Building Scale-out Production Applications Easy and Fast
• Robust Speech Recognition via Large-Scale Weak Supervision
• Large-scale Contrastive Language-Audio Pretraining with Feature Fusion and Keyword-to-Caption Augmentation
• Scaling Instruction-Finetuned Language Models
• LayoutLMv3: Pre-training for Document AI with Unified Text and Image Masking
• Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python
• Large-Scale Chemical Language Representations Capture Molecular Structure and Properties
• Longformer: The Long-Document Transformer
• MLIR: A Compiler Infrastructure for the End of Moore's Law
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension
• TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning
• AutoGraph: Imperative-style Coding with Graph-based Performance
• Glow: Graph Lowering Compiler Techniques for Neural Networks

Cited by

• Prompt Less, Smile More: MTP with Semantic Engineering in Lieu of Prompt Engineering

Discussions

• GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2 [hn, 3 points, 0 comments]

Related

• PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch
• Improving the GP 2 Compiler
• Compilation Techniques for Graph Algorithms on GPUs
• The Improved GP 2 Compiler
• Loo.py: transformation-based code generation for GPUs and CPUs