GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2

• URL: http://arxiv.org/abs/2509.16248v1
• Date: Wed, 17 Sep 2025 17:15:35 GMT
• Title: GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2
• Abstract summary: GraphMend is a high-level compiler that eliminates FX graph breaks in PyTorch 2 programs.<n>On NVIDIA 3090 and A40, GraphMend achieves up to 75% latency reductions and up to 8% higher end-to-end throughput.
• Score: 3.3228118981378816
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: This paper presents GraphMend, a high-level compiler that eliminates FX graph breaks in PyTorch 2 programs. Although PyTorch 2 introduced TorchDynamo and TorchInductor to enable just-in-time graph compilation, unresolved dynamic control flow and unsupported Python constructs often fragment models into multiple FX graphs. These fragments force frequent fallbacks to eager mode, incur costly CPU-to-GPU synchronizations, and reduce optimization opportunities. GraphMend addresses this limitation by analyzing and transforming source code before execution. Built on the Jac compilation framework, GraphMend introduces two code transformations that remove graph breaks due to dynamic control flow and Python I/O functions. This design allows PyTorch's compilation pipeline to capture larger, uninterrupted FX graphs without requiring manual refactoring by developers. Evaluation across eight Hugging Face models shows that GraphMend removes all fixable graph breaks due to dynamic control flow and Python I/O functions, driving the break count to 0 in 6 models and reducing it from 5 to 2 in another model. On NVIDIA RTX 3090 and A40 GPUs, GraphMend achieves up to 75% latency reductions and up to 8% higher end-to-end throughput. These results demonstrate that high-level code transformation is an effective complement to PyTorch's dynamic JIT compilation pipeline, substantially improving both usability and performance.

Related papers

• GTA: Graph Theory Agent and Benchmark for Algorithmic Graph Reasoning with LLMs [50.68372832207766]
  Graph Theory Bench (GT Bench) is a benchmark covering 24 classical graph problems in 44 task-structure settings.<n>GT Bench shows that accuracy is strongly tied to the input representation.<n>We propose the Graph Theory Agent (GTA), which pairs a preference-trained representation selector with plan-and-decompose scaffolding.
  arXiv  Detail & Related papers  (2026-09-10T22:50:10Z)
• VecFormer: Towards Efficient and Generalizable Graph Transformer with Graph Token Attention [61.96837866507746]
  VecFormer is an efficient and highly generalizable model for node classification.<n>VecFormer outperforms the existing Graph Transformer in both performance and speed.
  arXiv  Detail & Related papers  (2026-02-23T09:10:39Z)
• PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch [1.2334708058524546]
  We introduce PyGraph, a novel approach to automatically harness the power of NVIDIA Graphs within PyTorch2.<n>We evaluate PyGraph across various machine learning benchmarks, demonstrating substantial performance improvements over PyTorch2.
  arXiv  Detail & Related papers  (2025-03-25T15:47:54Z)
• Flatten Graphs as Sequences: Transformers are Scalable Graph Generators [7.341041463243556]
  AutoGraph is a scalable autoregressive model for attributed graph generation using decoder-only transformers.<n>It achieves state-of-the-art performance on synthetic and molecular benchmarks, with up to 100x faster generation and 3x faster training than leading diffusion models.
  arXiv  Detail & Related papers  (2025-02-04T10:52:14Z)
• InstructG2I: Synthesizing Images from Multimodal Attributed Graphs [50.852150521561676]
  We propose a graph context-conditioned diffusion model called InstructG2I. InstructG2I first exploits the graph structure and multimodal information to conduct informative neighbor sampling. A Graph-QFormer encoder adaptively encodes the graph nodes into an auxiliary set of graph prompts to guide the denoising process.
  arXiv  Detail & Related papers  (2024-10-09T17:56:15Z)
• TorchGT: A Holistic System for Large-scale Graph Transformer Training [27.67050123844626]
  Graph Transformer is a new architecture that surpasses GNNs in graph learning. We propose TorchGT, the first efficient, scalable, and accurate graph transformer training system. TorchGT boosts training by up to 62.7x and supports graph sequence lengths of up to 1M.
  arXiv  Detail & Related papers  (2024-07-19T08:21:42Z)
• GraphStorm: all-in-one graph machine learning framework for industry applications [75.23076561638348]
  GraphStorm is an end-to-end solution for scalable graph construction, graph model training and inference. Every component in GraphStorm can operate on graphs with billions of nodes and can scale model training and inference to different hardware without changing any code. GraphStorm has been used and deployed for over a dozen billion-scale industry applications after its release in May 2023.
  arXiv  Detail & Related papers  (2024-06-10T04:56:16Z)
• SpikeGraphormer: A High-Performance Graph Transformer with Spiking Graph Attention [1.4126245676224705]
  Graph Transformers have emerged as a promising solution to alleviate the inherent limitations of Graph Neural Networks (GNNs) We propose a novel insight into integrating SNNs with Graph Transformers and design a Spiking Graph Attention (SGA) module. SpikeGraphormer consistently outperforms existing state-of-the-art approaches across various datasets.
  arXiv  Detail & Related papers  (2024-03-21T03:11:53Z)
• A Graph is Worth $K$ Words: Euclideanizing Graph using Pure Transformer [47.25114679486907]
  We introduce GraphsGPT, featuring a Graph2Seq encoder that transforms Non-Euclidean graphs into learnable Graph Words. A GraphGPT decoder reconstructs the original graph from Graph Words to ensure information equivalence.
  arXiv  Detail & Related papers  (2024-02-04T12:29:40Z)
• PockEngine: Sparse and Efficient Fine-tuning in a Pocket [62.955793932377524]
  We introduce PockEngine: a tiny, sparse and efficient engine to enable fine-tuning on various edge devices. PockEngine supports sparse backpropagation and sparsely updates the model with measured memory saving and latency reduction. Remarkably, PockEngine enables fine-tuning LLaMav2-7B on NVIDIA Jetson AGX Orin at 550 tokens/s, 7.9$times$ faster than the PyTorch.
  arXiv  Detail & Related papers  (2023-10-26T19:46:11Z)
• Dynamic Graph Message Passing Networks for Visual Recognition [112.49513303433606]
  Modelling long-range dependencies is critical for scene understanding tasks in computer vision. A fully-connected graph is beneficial for such modelling, but its computational overhead is prohibitive. We propose a dynamic graph message passing network, that significantly reduces the computational complexity.
  arXiv  Detail & Related papers  (2022-09-20T14:41:37Z)
• DOTIN: Dropping Task-Irrelevant Nodes for GNNs [119.17997089267124]
  Recent graph learning approaches have introduced the pooling strategy to reduce the size of graphs for learning. We design a new approach called DOTIN (underlineDrunderlineopping underlineTask-underlineIrrelevant underlineNodes) to reduce the size of graphs. Our method speeds up GAT by about 50% on graph-level tasks including graph classification and graph edit distance.
  arXiv  Detail & Related papers  (2022-04-28T12:00:39Z)
• Boosting Graph Embedding on a Single GPU [3.093890460224435]
  We present GOSH, a GPU-based tool for embedding large-scale graphs with minimum hardware constraints. It employs a novel graph coarsening algorithm to enhance the impact of updates and minimize the work for embedding. It also incorporates a decomposition schema that enables any arbitrarily large graph to be embedded with a single GPU.
  arXiv  Detail & Related papers  (2021-10-19T15:25:04Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.