PyGraph让PyTorch更高效使用CUDA Graph，显著降低GPU启动延迟。
PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch
- 自动代码转换使机器学习程序适配CUDA Graph
- 消除图内核参数拷贝开销，提升执行效率
- 基于成本收益分析智能部署，适合大规模训练场景
机器学习工作负载每轮迭代需启动数百至数千个短时运行的GPU内核。随着GPU计算吞吐量快速提升，内核的CPU侧启动延迟正成为瓶颈。CUDA Graph通过单次调度重放一组内核，可消除每个内核的启动开销。然而，正确高效地部署CUDA Graph仍极困难。本文提出PyGraph——一个编译器框架，旨在最大化机器学习工作负载对CUDA Graph的覆盖与收益。它引入三项新优化：自动代码变换使ML应用适配CUDA Graph；消除图内核执行时的参数拷贝开销；基于成本收益分析选择性部署CUDA Graph。在来自TorchBench、HuggingFace和TIMM的25个机器学习工作负载上，PyGraph带来的收益比当前最主流的ML编译器PyTorch2高出一倍以上。PyGraph基于PyTorch2的编译框架构建，无需程序员干预。
原文摘要 · Abstract (English)
Machine learning (ML) workloads launch hundreds to thousands of short-running GPU kernels per iteration. With GPU compute throughput growing rapidly, CPU-side launch latency of kernels is emerging as a bottleneck. CUDA Graphs promise to address this by replaying a set of kernels with a single dispatch of the graph, removing per-kernel launch costs. However, CUDA Graphs remain surprisingly difficult to deploy correctly and efficiently. We present PyGraph - a compiler framework to maximize the coverage and benefits of CUDA Graphs for ML workloads. It introduces three novel optimizations: it applies automatic code transformations to make ML applications amenable to CUDA Graphs; it eliminates the parameter copy overheads for kernels executing in CUDA Graphs, and it selectively deploys CUDA Graphs guided by a cost-benefit analysis. For 25 ML workloads from TorchBench, HuggingFace, and TIMM, PyGraph more than doubles the benefit from deploying CUDA Graph compared to the most popular and widely used ML compiler, PyTorch2. PyGraph is built atop PyTorch2's compilation framework and requires no programmer intervention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。