This post explores the newly released PyTorch 2.0, which followed PyTorch 1.13.

OVERVIEW

I have used both PyTorch and TensorFlow, and if I had to compare them, I find PyTorch easier to use than TensorFlow. It is user-friendly—the process from defining classes to implementing training functions is very clean—and, more importantly, I have simply used it more. Any claim that one framework is better is entirely subjective, of course, and many deep-learning frameworks exist beyond PyTorch and TensorFlow.
Unlike TensorFlow, Meta (formerly Facebook), which developed PyTorch, has maintained remarkably consistent syntax and usage. As discussed later under Motivation, PyTorch's developers made building a platform centered on ease of use their highest priority. Moving beyond the previous generation, they made the major decision to rename PyTorch 1.x as PyTorch 2.0 in preparation for next-generation PyTorch.
One major advantage of PyTorch is how naturally it works with Python APIs, and PyTorch 2.0 promises faster training on top of that. The most important change in PyTorch 2.x is probably model.compile, discussed below.

What Is model.compile in PyTorch 2.0?

TensorFlow users will probably recognize the phrase. In TensorFlow, model.compile configures the optimizer, loss, and metrics used for fitting through method arguments. The loss is the objective function optimized directly, while a metric can be something such as classification accuracy.
model.compile in PyTorch 2.x is somewhat different. It preserves existing training code while allowing PyTorch to move beyond its traditional C++ base and operate at the Python level. Another advantage is that it is completely optional: choosing not to use features such as model.compile does not mean that you cannot use PyTorch 2.x. All code from earlier PyTorch versions remains compatible with PyTorch 2.x as-is. The technologies underlying torch.compile are as follows.

• TorchDynamo: Uses Python frame-evaluation hooks to improve the reliability of PyTorch programs. I do not yet understand every detail, but it appears to help capture graphs for operations such as backpropagation.
• AOTAutograd: Overloads PyTorch's existing Autograd engine. It traces automatic differentiation to produce a precomputed backward graph.
• PrimTorch: Canonicalizes roughly $2,000$ PyTorch operations into $250$ primitive operations, making it possible to build a complete PyTorch backend.
• TorchInductor: A deep-learning compiler that generates fast code for multiple accelerators and backends. For NVIDIA GPUs, it uses OpenAI Triton as a key component.

That is a rough translation. The feature is still close to a demo, so performance seems reliably improved on newer high-end GPUs but not yet meaningfully better on older models. I have a vague sense it may still be a little faster. The major benefit, in any case, is that existing PyTorch syntax works unchanged. It is hard to grasp the concept from these translated descriptions alone—even in English—so I will return to the details later.

Validation in PyTorch

Claims of speed need empirical validation, so the PyTorch team ran experiments. Because the only required change is to compile the model, they tested 163 open-source models: 46 Hugging Face Transformer models, 61 TIMM models, and 56 TorchBench models. These cover a broad range of deep-learning networks, from image classification to NLP and reinforcement learning. The goal was to show speed improvements across tasks.

The result is shown above. Without changing an open-source model's architecture or code, the team simply wrapped it with torch.compile and measured speedup and validation accuracy. Because speedup varies with data type, they measured both float32 and Automatic Mixed Precision (AMP), combining them as $0.75 \times AMP + 0.25 \times float32$. AMP received more weight because it is used more often in practice.
Across the 163 open-source models, torch.compile worked on 93%, and model training was 43% faster. These results were measured on an A100 server GPU. The team noted that desktop GPUs such as the 3090 might see smaller gains or even run more slowly.

Caveats: On a desktop-class GPU such as a NVIDIA 3090, we’ve measured that speedups are lower than on server-class GPUs such as A100. As of today, our default backend TorchInductor supports CPUs and NVIDIA Volta and Ampere GPUs. It does not (yet) support other GPUs, xPUs or older NVIDIA GPUs.

That is what they actually wrote. I installed PyTorch 2.0 full of anticipation, so learning that it did not work on my computer was deeply disappointing. I will wait patiently in the hope that I can use it someday.

Direction of PyTorch 2.x

The issues above probably arose because I tried PyTorch 2.0 early in its release, before everything was fully established. Returning to the PyTorch website after some time, I found that much had been updated. I am adding this section in the hope that a more stable PyTorch 2.0 is now available.

As PyTorch moved into the 2.x series, its developers began focusing on the compile method. They appear to plan to strengthen and gradually scale its ability to optimize models and accelerate training using the underlying technologies discussed above. Unlike earlier PyTorch releases, the PyTorch 2.0 series seems intended to evolve through ongoing research and interaction with users, using the resulting insights to provide better functionality in later versions.

The release includes the following remarks from well-known PyTorch users.

Sylvain Gugger, the primary maintainer of Hugging Face Transformers:

"With just one line of code to add, PyTorch 2.0 gives a speedup between 1.5x and 2.x in training Transformers models. This is the most exciting thing since mixed precision training was introduced!"

In short, applying model.compile to Transformers increased training speed by roughly 1.5–2.0×.

Ross Wightman the primary maintainer of TIMM (one of the largest vision model hubs within the PyTorch ecosystem):

“It just works out of the box with majority of TIMM models for inference and train workloads with no code changes”

This appears to mean that most built-in TIMM models became faster for inference and training immediately, without code changes. Did they really?

Luca Antiga the CTO of Lightning AI and one of the primary maintainers of PyTorch Lightning

“PyTorch 2.0 embodies the future of deep learning frameworks. The possibility to capture a PyTorch program with effectively no user intervention and get massive on-device speedups and program manipulation out of the box unlocks a whole new dimension for AI developers.”

The ability to use PyTorch efficiently without dissecting the program seems to imply that it can operate efficiently on a wide range of embedded platforms.

Motivation

PyTorch's developers placed the highest priority on preserving flexibility and hackability. They believed the system should provide users with maximum freedom, followed by performance improvements through platform optimization.

PyTorch launched in 2017. Since then, parallel processors and hardware accelerators such as GPUs have improved substantially in both memory-access speed and compute speed. To improve eager-execution performance, PyTorch naturally moved much of its code to C++—most PyTorch source has a C++ foundation—but this created a barrier that reduced users' ability to contribute code and thus its hackability. Eager execution is an imperative programming environment that executes operations immediately without first constructing a graph. TensorFlow users will probably recognize the concept. In this respect, PyTorch and TensorFlow pursued different directions in building deep-learning development environments.

The team ultimately concluded that improving eager-execution performance alone had limits. As early as July 2017, it decided to build a compiler for PyTorch. The compiler's goals were to make PyTorch faster without damaging the PyTorch experience—it could not become complicated or restrictive. Supporting dynamic shapes and programs was a central criterion for maintaining flexibility, user freedom, and ease of use.

The PyTorch 2.0 Compiler

As noted above, PyTorch has pursued compiler research since 2017. It divides the compiler into three components, of which graph acquisition is said to be the most difficult.

• Graph acquisition
• Graph lowering
• Graph compilation

After launch, PyTorch developed model-optimization tools including torch.jit.trace, TorchScript, FX tracing, and Lazy Tensors. None gave the developers the feel of the compiler they wanted. Some offered freedom at the cost of speed; greater speed reduced freedom. I also remember trying several scripts to use models trained on other platforms more efficiently and compatibly, only to encounter considerable inconvenience. TorchScript performed reasonably well, but depended heavily on how the code was written and required users to revise large amounts of it. Consequently, people less familiar with PyTorch tended not to use it.

The figure shows the role of each compiler technology. TorchDynamo and AOTAutograd first turn an eager-execution model into a graph and enable gradient computation. Operation-simplification methods such as Prim then replace PyTorch's many operations with simpler ones and simplify the graph. Finally, TorchInductor compiles the simplified graph.

TorchDynamo: Acquiring Graphs reliably and fast

TorchDynamo uses the CPython feature introduced as the “frame evaluation API” in PEP 523. Reliable compilation requires graph acquisition, and TorchDynamo can be viewed as a C-based framework that supports it. To test its effectiveness, the team ran more than $7,000$ PyTorch GitHub projects. TorchScript and other methods captured graphs only about $50\%$ of the time and imposed substantial code-modification overhead, whereas TorchDynamo captured graphs $99\%$ of the time without requiring code changes. It delivered both the flexibility and speed the developers wanted.

TorchInductor: fast codegen using a define-by-run IR

The growing adoption of the Triton language, which accelerates GPU operations, inspired the design of PyTorch 2.0's new backend compiler. The developers created a compiler backend that retains PyTorch eager execution and preserves PyTorch's existing characteristics.

Unlike TensorFlow, which defines a computation structure before evaluating it, PyTorch uses define-by-run, in which structure and computation proceed together. TorchInductor uses a training loop at the IR level and maps it to Triton code on GPUs or C++/OpenMP code on CPUs. The core loop contains few operators at the IR level and, more importantly, is implemented in Python, making it more hackable and extensible than C++.

AOTAutograd: reusing Autograd for ahead-of-time graphs

One goal of PyTorch 2.0 was faster training. Capturing only user-level code is not enough; backpropagation, which strongly affects training speed, must also be captured. The word capture appears frequently here. Think of taking a screenshot of something interesting or worth recording on a phone. If the user-defined sequence of model operations and the backpropagation obtained by differentiating the loss remain available throughout the training loop, they act as a guidebook that avoids redundant, cumbersome computation. The central idea was to reuse PyTorch's existing Autograd system. Using torch_dispatch, AOTAutograd improves speed by having the Autograd engine capture the backward pass in advance—ahead of time.

PrimTorch: Stable Primitive operators

Writing a PyTorch backend was harder than expected because PyTorch supports an enormous number of operators—more than $2,000$ by a detailed count.

Building a backend that handles the characteristics of every individual operator is exceptionally difficult. The PrimTorch project therefore created a smaller, stable operator set capable of expressing the many existing operators. Every PyTorch operator can be composed from a small operator set. The Prim project defines two such sets.

• Prim operators ($\sim 250$): Mostly low-level operators suited to the compiler level.
• ATen operators ($\sim 750$): Operations suited to the ATen level; higher-level than Prim operators.

Installing PyTorch 2.0

First, check your GPU specifications.

nvidia-smi

My CUDA version was 11.7, so I installed the latest nightly build for that specification.

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117

For CUDA 11.6, use the following installation command.

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu116

If you have only a CPU, use this command instead.

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cpu

Using PyTorch 2.0

import torch
import torchvision.models as models

model = models.resnet18().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
compiled_model = torch.compile(model)

x = torch.randn(16, 3, 224, 224).cuda()
optimizer.zero_grad()
out = compiled_model(x)
out.sum().backward()
optimizer.step()

compiled_model accelerates the user-defined network by optimizing its forward method. In more detail, the compilation API is as follows.

def torch.compile(model: Callable,
  *,
  mode: Optional[str] = "default",
  dynamic: bool = False,
  fullgraph:bool = False,
  backend: Union[str, Callable] = "inductor",
  # advanced backend options go here as kwargs
  **kwargs
) -> torch._dynamo.NNOptimizedModule

If no option other than the model is supplied, compilation applies the default settings shown above. Let us examine each argument.

• mode The mode specifies how compilation should optimize the model. Default mode compiles efficiently without taking too long or consuming excessive memory. reduce-overhead uses somewhat more memory to reduce framework overhead, while max-autotune spends a long time compiling code intended to train as quickly as possible.

# API NOT FINAL
# default: optimizes for large models, low compile-time
#          and no extra memory usage
torch.compile(model)

# reduce-overhead: optimizes to reduce the framework overhead
#                and uses some extra memory. Helps speed up small models
torch.compile(model, mode="reduce-overhead")

# max-autotune: optimizes to produce the fastest model,
#               but takes a very long time to compile
torch.compile(model, mode="max-autotune")

• dynamic dynamic is a Boolean that controls whether to enable a code path for dynamic shapes. Compiler optimization can sometimes make a program incompatible with dynamic shapes; this setting lets users control compilation accordingly. I do not yet understand it fully, but it appears to allow the compiler to compile graphs flexibly when data shapes change.

• fullgraph fullgraph resembles Numba's nopython. It compiles the entire program into a single graph and, if that fails, displays an error explaining why. It is optional.

• backend backend selects the compiler backend. The default is TorchInductor, described above, though other options are available. I have not investigated them in detail.

Features Available After Compilation

Reading and updating Attributes

One of eager mode's most convenient features is the ability to access a model's weights or read their values during training—for example, model.conv1.weight. TorchDynamo preserves this behavior and automatically recompiles the relevant section when it detects that an attribute has changed.

# optimized_model works similar to model, feel free to access its attributes and modify them
optimized_model.conv1.weight.fill_(0.01)

# this change is reflected in model

Serialization

You can serialize the state dict of either the optimized model or the original model. Because both refer to the same parameters, the following two lines produce the same result.

torch.save(optimized_model.state_dict(), "foo.pt")
# both these lines of code do the same thing
torch.save(model.state_dict(), "foo.pt")

Serializing the model object itself is different. To avoid an error, save the original model, not the optimized model.

torch.save(optimized_model, "foo.pt") # Error
torch.save(model, "foo.pt")           # Works

Keep this distinction in mind when saving.

Model inference

For model inference, performing a warm-up step after creating a compiled model with compile reportedly reduces initial latency. This still appears to be an area for improvement in PyTorch 2.x. A future export method is intended to stabilize even the environmental variables affecting latency. I will revisit the topic as it evolves.

# API Not Final
exported_model = torch._dynamo.export(model, input)
torch.save(exported_model, "foo.pt")

Debugging

Compiled programs are difficult to debug. A program may conflict with compile mode; compiled operations may not match eager-mode operations exactly; or compilation may fail to deliver the expected speedup. If compiled code conflicts with the program or its results differ from eager mode beyond the expected error tolerance, the user's code is not necessarily at fault. PyTorch provides a minifier to help users debug and reproduce such issues. It automatically reduces a problem to a small code snippet, which can reproduce the failure and be submitted in a GitHub issue. This lets the PyTorch team identify the cause quickly without inspecting an enormous codebase. For speedup problems, torch._dynamo.explain can identify performance bottlenecks known as graph breaks.

Distributed

torch.distributed is one way to use multiple GPUs for parallel computation. PyTorch's two representative distributed wrappers, DDP (DistributedDataParallel) and FSDP (FullySharedDataParallel), both work correctly after compilation and reportedly improve performance and memory efficiency over eager mode.

Summary

I still do not understand every detail, but in summary, compile optimizes existing PyTorch models that run in eager mode while preserving their familiar convenience: easy access to model parameters, parallel computation, and freedom in structuring code. PyTorch 2.x will likely make compile increasingly convenient and extensible. Compared with the more restrictive and complex TensorFlow, it may continue to establish itself as a more widely adopted framework while evolving to meet developers' demand for high performance.