arXiv:2509.16248cs.PLcs.LG2025-09被引 1

自动修复PyTorch 2中的图断裂问题,提升训练速度与可用性。

GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2

  • 通过源码分析识别图断裂模式并安全转换代码
  • 在27个模型中消除73%的图断裂,最高提速26倍
  • 适合需要性能优化的深度学习开发者使用

本文提出GraphMend,一种编译器技术,可自动修复PyTorch 2程序中的FX图断裂。尽管PyTorch 2引入TorchDynamo和TorchInductor实现即时图编译,但部分代码模式仍导致图断裂,迫使执行回落至Python急模式,引发昂贵的CPU-GPU同步并减少优化机会。我们对195个Hugging Face模型的调查显示,13.8%的模型存在图断裂。GraphMend通过源码级程序分析与转换,自动消除可修复的断裂。它分析AST级程序结构以识别图断裂模式,并仅在语义可静态保证的前提下应用转换。该方法使PyTorch能捕获更大、连续的FX图,无需开发者手动重构。我们在所有27个存在图断裂的模型上评估GraphMend,成功消除107/147个断裂(73%),并完全修复21个模型的全部断裂。在NVIDIA GPU实验中,GraphMend实现最高26倍冷启动加速、平均5倍加速,以及最高1.39倍稳定状态前向传播加速。结果表明,语义感知的源码级分析与转换是PyTorch动态JIT编译管道的有效补充,显著提升可用性与性能。

原文摘要 · Abstract (English)

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.

PyTorch图优化编译器性能提升

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。