PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation

By Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, D. Berard, Evgeni Burovski and 39 more. International Conference on Architectural Support for Programming Languages and Operating Systems, 2024

1,268 citations recorded by the source corpus.

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.

Subfields: Hardware and Architecture, Information Systems, Software

View Paper: https://www.semanticscholar.org/paper/a509670659e8b054e2b7d1b6f8a0bc722398fa62

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Hardware and Architecture

1. Shoufa Chen
2. Yi Jiang
3. Peize Sun
4. Bingyue Peng
5. Shilong Zhang
6. Ping Luo
7. Zehuan Yuan

In Information Systems

1. Ping Luo
2. Faisal Ladhak

Citations in the corpus

• Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation (2024, methodology): “All models are implemented with PyTorch 2 [Ansel et al. 2024] and trained on 80GB A100 GPUs.”
• Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference (2024, positive): “Mod-ernBERT uses a mixture of Flash Attention 3 for global attention layers and Flash Attention 2 (Dao, 2023) for local attention layers. torch.compile We leverage PyTorch’s built-in compiling (Ansel et al., 2024) to improve the training efficiency by compiling all compatible modules.”