PyTorch previews 2.0 with an optional compiler

PyTorch previewed version 2.0 on 2 December 2022, adding an optional compiler through torch.compile while retaining its usual Python workflow.

The feature is available in nightly builds. The announcement-day overview targets the first stable release for early March 2023.

Capture first, then compile

TorchDynamo captures graphs of operations from Python bytecode. When it encounters unsupported Python constructs, it can split the graph and let those portions continue in eager execution. This approach gives developers a way to keep existing code running while compiling the parts the system can handle.

That distinction matters when testing a model: successful execution does not mean the compiler captured the entire program. The earlier developer report separates its tests with fallback enabled from the more restrictive requirement to capture a complete graph.

TorchInductor provides the code-generation side. Its Python implementation targets OpenAI’s Triton for NVIDIA GPUs and C++ with OpenMP for CPUs. The design follows PyTorch’s handling of tensors, including views and mutations, rather than requiring every model to fit a different execution model.

Captured graphs reach device backends while unsupported portions use eager execution. — Simplified compiler routes, with eager fallback.

A preview to evaluate

PyTorch reports 43% faster training on an A100 across its 163-model evaluation, using an average weighted toward mixed precision. Compilation worked in 93% of cases, with lower gains on an RTX 3090. These are the project’s measurements.

The December preview has limits: its default TorchInductor backend supports CPUs and NVIDIA Volta and Ampere GPUs. Dynamic-shape support remains early.

Compile time also makes an initial run slower. Warm-up runs belong in performance testing before judging whether the compiled model helps a particular workload.