GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2

arXiv:2509.16248 · cs.PL, cs.LG, cs.SE · Submitted 2025-09-17 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2".

Jane: GraphMend addresses a critical performance bottleneck in PyTorch 2 where specific code patterns cause "graph breaks," forcing the execution pipeline to fall back to costly Python eager mode.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, what we’re seeing here is the paper titled "GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch two" and the authors are a team from the University of Michigan tackling this issue head-on. It’s interesting to see how many experts it takes to tackle a problem that seems common across production models like T5-small and Phi-four-mini-instruct.

Jane: They introduced GraphMend as a compiler technique that automatically fixes those graph breaks, which are the things that force the system to fall back into slow Python eager mode execution. Basically, they want to make the dynamic JIT compilation pipeline work more consistently across different code structures.

Lu: The core of what they did is applying source-level program analysis and transformations to identify and eliminate those breaks before bytecode is even generated, which lets TorchDynamo capture a single, continuous FX graph instead of multiple disjoint ones.

Meng: That focus on the source level analysis is what’s compelling because it means the fix happens before the code gets compiled into bytecode, which should save us a lot of debugging time later on when we encounter these issues in our own proprietary code.

Lalam: For me, this approach makes sense because it suggests that we can address these compilation hurdles at a level where the meaning of the computation is still intact, rather than just patching runtime errors after they happen.

The paper's summary: Tom: To get into the summary of "GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch two" it explains that the system automatically fixes these graph breaks by analyzing the program structure at the Abstract Syntax Tree level and applying transformations only when their semantic preservation can be statically established. This lets TorchDynamo capture a single, continuous FX graph instead of multiple disjoint ones.

Jane: In simpler terms, they are taking the source code before it becomes machine instructions and making smart changes to it so the compiler doesn't get confused by things like data-dependent control flow or I/O calls that normally cause a break. They are rewriting parts of the code to ensure everything stays within one unified structure.

Lu: The paper outlines three specific, semantically-preserving AST-level transformations they use to solve this problem: rewriting data-dependent branches into `torch.where` expressions, buffering side effects like print calls using Graph-Epilogue Deferred Side Effects, and replacing certain error patterns with graph-native assertions.

Meng: The focus on source-level analysis is what’s compelling here because it means the fix happens before the code even gets compiled into bytecode, which should save us a lot of debugging time later on when we encounter these issues in our own proprietary code.

Lalam: For me, I think the idea of deferring side effects until after the main computation is very important; it means I can focus my processing resources entirely on the core task without worrying about external logging or printing interrupting that flow.

The paper's improvements: Tom: So, what we’re seeing in terms of actual results with "GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch two" is that they successfully eliminated one hundred seven out of one hundred forty-seven total graph breaks across the models they tested, which is a fix rate of about seventy-three percent. They also noted that they fully fixed all breaks in twenty-one models out of the ones studied.

Jane: Those performance gains are quite substantial when you look at the GPU results; they reported up to a twenty-six times cold-start speedup, averaging about five times faster, and a steady-state forward pass speedup of up to one point three nine times. That kind of acceleration really impacts how fast we can get results from the AI models.

Lu: The improvements are substantial because they manage to fix structural issues that were previously considered inherent limitations of the PyTorch two compilation pipeline itself, proving that these breaks aren't just developer errors but are a structural challenge in production code.

Meng: The fact that they managed this with only source-level analysis and transformations suggests we don't need to fundamentally redesign the way we write our models, which is a huge practical win for engineers trying to keep up with the latest compiler requirements.

Lalam: I think the ability to fix these breaks automatically means that I can be deployed in environments where performance is highly sensitive, because my execution path won't be interrupted by these synchronization pauses.

Conclusion: Tom: To wrap up on "GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch two" the main implication is that we can significantly reduce the compilation overhead and synchronization bottlenecks that plague large model execution when using TorchDynamo. It gives us a much smoother path from code to high-performance GPU execution.

Jane: Exactly, Tom; it moves the burden of fixing these performance hurdles from manual refactoring to an automated compiler pipeline that understands how to preserve the program's meaning while enabling continuous graph capture. It really makes deploying complex models much more practical.

Lu: From a research perspective, the study confirms that high-level source code patterns are indeed a major cause of these breaks, validating the need for tools like GraphMend to analyze AST and CFG structures. This points toward better static analysis techniques being key for future compiler improvements.

Meng: I think the practical impact here is huge because it directly translates to faster time-to-first-token, which is a major bottleneck in serverless inference environments we are all trying to optimize.

Lalam: My perspective is that this advancement in fixing graph breaks allows for much more robust and consistent AI systems overall, because the underlying computation runs with less interruption, which should lead to more reliable and efficient interactions for everyone using these models.

University of Michigan · Jaseci Labs

cs.PL, cs.LG, cs.SE

Submitted: 2025-09-17

Updated: 2026-09-03

Comments: Accepted for publication at the International Symposium on Code Generation and Optimization (CGO 2027)

Code: https://github.com/Jaseci-Labs/jaseci

Project page: https://sglproject.github.io/advanced_features/server_arguments.html

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 79/100

The gist: GraphMend addresses a critical performance bottleneck in PyTorch 2 where specific code patterns cause "graph breaks," forcing the execution pipeline to fall back to costly Python eager mode.

Key concepts

Graph Breaks
These are specific code patterns in PyTorch 2 that cause the execution pipeline to break, forcing the system to fall back into slow Python eager mode instead of efficient dynamic JIT compilation.
GraphMend
A compiler technique introduced by the University of Michigan team that automatically fixes graph breaks. It works by applying source-level program analysis and transformations at the Abstract Syntax Tree level before bytecode generation.
Source-Level Analysis
The process of analyzing the program structure directly from the source code, rather than after it is compiled into machine instructions. This allows GraphMend to make fixes before compilation begins, saving debugging time.
AST-level Transformations
Specific changes applied to the Abstract Syntax Tree during GraphMend's process. These include rewriting data-dependent branches into torch.where expressions and buffering side effects like print calls.

Terminology

Summary

GraphMend addresses a critical performance bottleneck in PyTorch 2 where specific code patterns cause graph breaks, forcing the execution pipeline to fall back to costly Python eager mode. This fragmentation prevents the compiler from capturing a single, continuous FX graph, leading to CPU-GPU synchronization overhead and reduced optimization opportunities. Our investigation of 195 Hugging Face models revealed that while these issues are not due to developer inexperience but structural challenges in production code, GraphMend provides an automated solution, enabling PyTorch’s dynamic JIT compilation pipeline to capture larger, uninterrupted FX graphs without manual refactoring.

The Problem: Graph Breaks and System Overhead

When TorchDynamo attempts to symbolically evaluate a function for the FX graph capture process, it encounters unsupported Python operations—such as data-dependent control flow (e.g., if x.sum > 0), I/O calls, or runtime validation guards. These unsupported elements cause the compiler to insert graph breaks, resulting in multiple disjoint FX graphs instead of one unified structure. The consequences are severe:

  • The regions between these graphs execute in PyTorch eager mode, returning control to the Python interpreter.

  • This introduces significant overhead from Python dispatch and forces CPU-GPU synchronization, especially when the break is caused by data-dependent checks.

  • Each disjoint FX graph must be compiled independently during the cold run, multiplying compilation overhead.

How GraphMend Works

GraphMend is a compiler technique that applies source-level program analysis and transformations to eliminate these breaks before bytecode generation. It extends the standard PyTorch 2 compilation pipeline by integrating with the Jaseci framework, which accepts Python source code. GraphMend uses a conservative legality analysis over the Abstract Syntax Tree (AST), Control Flow Graph (CFG), and symbol table, applying transformations only when semantic preservation can be statically established. This approach is necessary because high-level semantic information is lost when low-level bytecode is generated.

Key Transformation Rules

GraphMend automates the fix for graph breaks by applying three specific, semantics-preserving AST-level transformations:

  1. Predicated Data-Dependent Control Flow: It rewrites data-dependent branches into torch.where expressions, ensuring both paths remain within the single FX graph structure.

  2. Graph-Epilogue Deferred Side Effects: It buffers print and logger calls, deferring their execution until after the main graph computation is finished to prevent runtime interruption of the tracing process.

  3. Predicated Trap Lowering: It replaces if not condition: raise patterns with a graph-native assertion, torch. assert async, which maintains runtime safety checks without breaking tracing.

Evaluation and Performance Gains

We evaluated GraphMend on 27 models that exhibited graph breaks in our investigation. The results demonstrate significant effectiveness:

  • GraphMend successfully eliminated 107 of 147 total graph breaks (a fix rate of 73%).

  • It fully fixed all breaks in 21 models, while the remaining unfixed issues—such as dynamic shape operators or tensor.item calls—are deemed untraceable at the source level.

  • In NVIDIA GPU experiments, GraphMend achieved substantial performance improvements: up to a 26× cold-start speedup (5x on average) and up to a 1.39× steady-state forward pass speedup, confirming that semantics-aware analysis is effective at improving both usability and performance.

Improvements for AI systems

Based on a thorough analysis of the GraphMend paper, here are the specific improvements and resulting capabilities for implementing AI systems using this technology:

GraphMend introduces automated, compiler-level optimizations that fundamentally eliminate common bottlenecks in PyTorch 2 execution. These improvements move the burden of complex code refactoring from the model developer to the automated compiler pipeline.

1. Automated Semantic Preservation:

  • Improvement: The system performs source-level program analysis using Abstract Syntax Trees (AST) and Control Flow Graphs (CFG). It automatically identifies and applies transformations only when their semantic preservation can be statically established. This ensures that models remain functionally correct while allowing the compiler to optimize.

  • Impact on Developers: Developers are relieved of manual, complex refactoring tasks—such as hoisting branches or restructuring functions to satisfy torch.compile requirements`.

2. Automated Graph Consolidation (Addressing Data-Dependent Control Flow):

  • Improvement: The system automatically converts data-dependent Python if/else blocks (where the condition relies on runtime tensor values) into equivalent, graph-native operations, specifically torch.where. This transformation ensures both execution paths are captured within a single continuous FX graph.

  • Impact on System Architecture: Enables the seamless integration of previously uncompilable models into fully optimized PyTorch 2 pipelines without requiring custom model architecture changes for torch.compile compatibility.

3. Automated Side-Effect Management:

  • Improvement: The system detects Python side effects (e.g., print or logging calls) that break the continuous execution stream. It automatically buffers these calls and defers their execution to occur after the entire graph computation is complete (Graph-Epilogue Deferred Side Effects).

  • Impact on Performance: Eliminates forced CPU-GPU synchronization events caused by runtime I/O, ensuring uninterrupted GPU utilization.

4. Automated Safety Guard Conversion (Predicated Trap Lowering):

  • Improvement: The system replaces Python validation guards (e.g., if not torch.equal(A, B): raise Error) with a graph-native assertion (torch. assert async). This maintains the program's runtime safety checks while ensuring the entire operation remains inside the FX graph.

  • Impact on Reliability: Allows high-level input validation and data integrity checks to be performed within a compiled, optimized context without sacrificing functional correctness.

By implementing GraphMend, an AI system is transformed from one constrained by manual refactoring limitations to one capable of peak performance and seamless deployment.

1. Significant Performance Gain (Cold Start):

  • The system enables the capture of a single, unified CUDA graph where previously multiple disjoint subgraphs existed. This drastically reduces the compilation overhead during the initial run (cold start).

  • Result: Up to 26 times cold-start speedup is achievable, which is critical for improving Time-to-First-Token in serverless or auto-scaling inference environments.

2. Peak Runtime Performance (Steady State):

By eliminating graph breaks, the system ensures the entire forward pass executes as a continuous CUDA region.

  • Result: Up to 1.39 times steady-state forward pass speedup is achieved by maximizing GPU utilization and enabling kernel fusion across what were previously broken execution boundaries.

3. Enhanced Serving Framework Compatibility:

The system successfully enables full graph capture for models that were previously incompatible.

  • Result: Models that failed to run in frameworks like vLLM or SGLang due to graph breaks can now pass the fullgraph=True requirement, dramatically expanding the scope of deployable, high-performance models.

4. Optimized Resource Utilization:

By removing CPU-GPU synchronization points and enabling continuous execution.

  • Result: The system optimizes resource scheduling and minimizes idle time on high-end GPUs (NVIDIA RTX 3090, A40, H100), maximizing the computational density of the inference workload.

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.

Sources

Related papers