一句话：PyTorch 2 的 torch.compile 借助动态字节码变换捕获计算图并编译，在真实模型上取得显著加速。
该论文介绍 PyTorch 2 中 torch.compile 功能的两个核心组件 TorchDynamo 与 TorchInductor。问题在于 eager 模式的 PyTorch 虽灵活却难以应用编译器优化。TorchDynamo 是一种 Python 级 JIT 编译器，在不牺牲 Python 灵活性的前提下，于执行前动态修改 Python 字节码，将 PyTorch 操作序列提取为 FX 图，再交由可扩展后端 JIT 编译。TorchInductor 是默认后端，将 PyTorch 程序翻译为面向 GPU 的 OpenAI Triton 和面向 CPU 的 C++。实验表明，TorchDynamo 比先前方法更稳健地捕获图且额外开销极小；在 NVIDIA A100 上对 180 多个真实模型，TorchInductor 实现推理 2.27 倍、训练 1.41 倍的几何平均加速，优于其他六个编译器。这为 eager 框架通过编译器应用优化提供了新途径。
This paper introduces two extensions to the popular PyTorch machine learning framework, TorchDynamo and TorchInductor, which implement the torch.compile feature released in PyTorch 2. TorchDynamo is a Python-level just-in-time (JIT) compiler that enables graph compilation in PyTorch programs without sacrificing the flexibility of Python. It achieves this by dynamically modifying Python bytecode before execution and extracting sequences of PyTorch operations into an FX graph, which is then JIT compiled using one of many extensible backends. TorchInductor is the default compiler backend for TorchDynamo, which translates PyTorch programs into OpenAI's Triton for GPUs and C++ for CPUs. Results show that TorchDynamo is able to capture graphs more robustly than prior approaches while adding minimal overhead, and TorchInductor is able to provide a 2.27× inference and 1.41× training geometric mean speedup on an NVIDIA A100 GPU across 180+ real-world models, which outperforms six other compilers. These extensions provide a new way to apply optimizations through compilers in eager mode frameworks like PyTorch.
@inproceedings{AnselYHGJVBBBBC24,
title = {PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation},
author = {Jason Ansel and Edward Yang and Horace He and Natalia Gimelshein and Animesh Jain and Michael Voznesensky and Bin Bao and Peter Bell and David Berard and Evgeni Burovski and Geeta Chauhan and Anjali Chourdia and Will Constable and Alban Desmaison and Zachary DeVito and Elias Ellison and Will Feng and Jiong Gong and Michael Gschwind and Brian Hirsh and Sherlock Huang and Kshiteej Kalambarkar and Laurent Kirsch and Michael Lazos and Mario Lezcano and Yanbo Liang and Jason Liang and Yinghai Lu and C. K. Luk and Bert Maher and Yunjie Pan and Christian Puhrsch and Matthias Reso and Mark Saroufim and Marcos Yukio Siraichi and Helen Suk and Shunting Zhang and Michael Suo and Phil Tillet and Xu Zhao and Eikan Wang and Keren Zhou and Richard Zou and Xiaodong Wang and Ajit Mathews and William Wen and Gregory Chanan and Peng Wu and Soumith Chintala},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620665.3640366}
}