PyTorch 2.0 goes live with torch.compile and faster Transformers

PyTorch 2.0 is here, and it brings the biggest shift to the machine learning framework in years. It keeps the eager-mode Python experience developers love, but completely redesigns the compiler under the hood so you can speed up training and inference without rewriting your codebase.

The headline feature is torch.compile, an additive and optional API backed by 4,541 commits and 428 contributors since version 1.13.1.

What it does and how it works

Wrapping an existing model with model = torch.compile(model) now produces an optimized binary automatically. Underneath, PyTorch deploys four new components: TorchDynamo safely intercepts execution graphs, AOTAutograd traces backward passes ahead of time, PrimTorch shrinks over 2,000 operators down to roughly 250 primitives, and TorchInductor generates machine code.

For Nvidia and AMD hardware, TorchInductor relies on OpenAI Triton to compile kernels that rival hand-tuned CUDA libraries like cuBLAS. Across 165 open-source benchmarks, PyTorch found that torch.compile worked out of the box on 93% of models, boosting performance by an average of 20% on float32 and 36% under AMP mixed precision.

Transformers also get first-class treatment. The team embedded scaled dot product attention directly into torch.nn.functional. That means you get battle-tested implementations like FlashAttention and xFormers natively, without installing flaky third-party extensions.

Elsewhere, Apple Silicon users get a vastly improved MPS backend supporting over 300 operators, while AWS tuned Graviton3 CPU performance on models like BERT and ResNet-50.

How to get started

PyTorch 2.0 is accessible through the standard PyTorch wheels via pip or conda, fully documented on the official get-started site.

You do not need to rewrite your training scripts. Simply import torch, define your module, wrap it with torch.compile, and run your standard forward and backward steps.

Mac users planning to train on Apple Silicon should update to the latest macOS release to take advantage of the expanded MPS operator coverage and correctness fixes.

Why it matters

If you train LLMs or fine-tune vision models, PyTorch 2.0 is an immediate win. It gives developers the speed advantages typically associated with JAX or custom C++ backends, while letting you keep the Pythonic syntax everyone knows.

Keep in mind that torch.compile remains in Beta. While it hits 93% compatibility on popular benchmarks, dynamic edge cases might still fall back to eager mode. We suggest testing output tensors for numerical parity before swapping it straight into heavy production jobs.