PyTorch 2 Extensions for Dynamic Graph Compilation

This technical report presents two extensions to the PyTorch machine learning framework, namely TorchDynamo and TorchInductor, which implement the torch.compile feature introduced in PyTorch 2. TorchDynamo functions as a Python-level just-in-time (JIT) compiler that facilitates graph compilation in PyTorch programs while maintaining the flexibility inherent to Python. It achieves this by dynamically altering Python bytecode prior to execution and extracting sequences of PyTorch operations into an FX graph, which is subsequently JIT compiled using various extensible backends. TorchInductor serves as the default compiler backend for TorchDynamo, translating PyTorch programs into OpenAI’s Triton for GPUs and C++ for CPUs. The report details experimental results indicating that TorchDynamo captures graphs more effectively than previous methods while introducing minimal overhead. Additionally, TorchInductor demonstrates significant performance improvements, achieving a 2.27× speedup in inference and a 1.41× speedup in training across a range of real-world models, thus providing a novel approach to optimization in eager mode frameworks like PyTorch.