PyTorch 2.0 and torch.compile: Graph Compilation Without Rewriting Code

PyTorch always had a great advantage: code runs line by line like normal Python, making it very easy to debug. The problem is that this approach leaves a lot of performance on the table compared to frameworks that compile everything before executing.

PyTorch 2.0 introduces torch.compile, a single function you can apply to your model. Under the hood, a system called TorchDynamo analyzes your Python code on the fly, captures the operations graph, and passes it to an optimization backend called Inductor, which generates optimized kernel code for CPU or GPU.

The practical result is significant: on the same models and hardware, torch.compile achieves on average between 30% and 200% speedup over uncompiled code, without changing a single line of your model. For those training large transformers, this translates directly into fewer compute hours and lower costs. PyTorch can finally compete with JAX, which had this advantage all along, while maintaining the ease of use that made PyTorch dominant in research.