PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation
Jason Ansel (Meta AI)
Research Seminar
Thursday, November 7, 2024, 3:30 pm
Abstract
PyTorch 2 leverages compilers to deliver faster training and inference without sacrificing the usability and flexibility for which PyTorch is known. It is fully backward compatible and continues to provide an imperative, extensible, easy-to-debug, and Pythonic programming environment for AI researchers, data scientists, and engineers. PyTorch 2 achieves a 2.27x speedup in inference and a 1.41x speedup in training (geometric mean) on an NVIDIA A100 GPU across 180+ real-world models, outperforming six other compilers.
This talk will explore the key technologies behind PyTorch 2: TorchDynamo and TorchInductor.
TorchDynamo is a Python-level JIT compiler designed to accelerate unmodified PyTorch programs. It hooks into the CPython frame evaluation API to dynamically modify Python bytecode just before execution. When rewriting bytecode, it extracts sequences of PyTorch operations into an FX Graph, which is then compiled with various extensible backends. TorchDynamo generates compiled graph fragments that can seamlessly mix with regular Python execution, giving better performance while maintaining the flexibility of Python.
TorchInductor, the default compiler backend in PyTorch 2, translates PyTorch programs into OpenAI's Triton for GPUs, OpenMP/C++ for CPUs, and now contains an experimental Halide backend targeting both. It maintains the flexibility and dynamism of PyTorch by employing
abstractions similar to those in PyTorch’s eager mode. TorchInductor introduces a new define-by-run loop-level intermediate representation
(IR), facilitating easy integration of new operator lowerings. Additionally, as it is implemented in Python, TorchInductor is straightforward for users to extend and modify to fit their specific needs.
Bio
Jason Ansel is a Principal Research Scientist at Meta AI and technical lead for PyTorch compilers. He started the TorchDynamo and TorchInductor projects, bringing flexible graph capture and a high-performance compiler to PyTorch 2. Jason holds a Ph.D. from MIT CSAIL, earned in 2014, with research focusing on the intersection of machine learning, compilers, and programming languages.
This talk is viewable on our YouTube channel.