torch.compile adds a compiled wrapper around a PyTorch model; Dynamo captures operations into an FX graph with guards, AOT Autograd prepares forward and backward graphs, and Inductor generates hardware-specific code — the speed-up depends on the model and workload, since compilation is an opportunity to fuse and optimize operations, not a guaranteed win.
AOT Autograd is a PyTorch compilation component used as the second stage of torch.compile. It prepares a captured graph for a backend by decomposing high-level operations and capturing the backward graphs needed for training.
Dynamo is the first compilation stage used by PyTorch's torch.compile system. It traces Python execution, captures tensor operations in an FX graph, and records guards—assumptions about the inputs—that determine when the captured graph can be reused.
Inductor is PyTorch's default backend compiler for optimizing compiled models. It fuses operations, schedules work, and generates Triton kernels for GPUs or C++ code for CPUs as part of the torch.compile pipeline, alongside TorchDynamo and AOT Autograd.
PyTorch is a Python machine-learning library for tensor computation and deep neural networks, with CPU and GPU execution. Its tensor library provides NumPy-like operations, while its tape-based autograd system uses reverse-mode automatic differentiation for differentiable tensor operations and dynamically defined networks. The project also includes neural-network, compilation, multiprocessing, and data-loading components, and supports extensions through Python packages such as NumPy, SciPy, and Cython.
torch.compile is PyTorch's model compilation tool. It inspects a model before execution and converts as much of its computation as possible into an optimized region, allowing compiler transformations such as kernel fusion, operation reordering, and memory planning. Python scalar reads such as `.item()` can cause graph breaks that fragment these regions and force synchronization between Python and the accelerator, reducing available optimization and potentially slowing execution.
Triton is a kernel language that PyTorch Inductor commonly generates for GPU targets. In the torch.compile pipeline, it is part of the stage that generates optimized code for PyTorch models.
Searchable transcript of What Is torch.compile? Dynamo, PyTorch & Faster AI Models Explained — IBM Technology (06:53). Search for a phrase, then click its timestamp to jump straight to that moment in the video.
Captions sourced from the original video on YouTube, published by IBM Technology. The video, its captions and all related intellectual property remain the property of their respective owners; AINotes claims no ownership. Provided for research, accessibility and search — see the Transcript Notice and Copyright Policy.
00:00 You have a PyTorch model that runs correctly and you want to make it faster. You add one line model equals torch.compile. The first run may take longer because PyTorch has to compile the code. After that, the model may run faster. But what is PyTorch actually doing? And why does compilation help some models much more than it helps others? To answer that, we need to look at the three stages.
00:19 Dynamo, AOT, Autograd, and inductor. In normal PyTorch execution, Python runs your model and dispatches operations as it reaches them. Some operations already call highly optimized libraries. Others may require separate kernel launches. A kernel is a piece of code that runs on hardware such as a GPU. For small operations, the overhead of launching work and moving intermediate data can become significant.
00:48 Torch compile gives PyTorch an opportunity to capture a section of the computation, optimize it and generate code that executes it more efficiently. The important word here is opportunity. Compilation does not guarantee a speed up and it does not necessarily turn and the entire model into a single graph. Here is the line you write model equals torch.compile.
01:14 This creates a compiled wrapper around the model you already have. With the standard setup, compilation is lazy. This means that PyTorch starts tracing and compiling when you first call the model with the inputs. That first call provides information such as tensor shapes, data types, and devices. PyTorch uses this information to specialize the generated code.
01:37 Now let's look at what happens during that first call. The first stage is Dynamo. Dynamo traces your Python execution and captures the tensor operations it can represent in an FX graph. An FX graph is a representation of the operations. Dynamo doesn't just trace the operations. It also records the information that flows between them. Rather than simply storing your original code, Python source code, Dynamo also records assumptions about the inputs and the program state.
02:05 These are called the guards. For example, a compiled program may be valid for a particular tensor shape or data type. On later calls, PyTorch checks whether those assumptions still hold before reusing the compiled code. If the assumptions no longer hold, PyTorch may decide it needs to compile a new version again. So, Dynamo's job is not just to record operations.
02:35 It determines what can be captured and under which conditions that captured code can be reused. The second stage is AOT autograd. Its role is to prepare the captured computation for the backend. This includes decomposing some higher level operations into lower level operations that the compiler can work with. It's better to think of this as preparing the graph, not simply making it smaller.
03:00 A decomposition, for example, can actually turn one operation into several. For training. AOT Autograd also captures the backward computation used to calculate the gradients. It can prepare forward and backward graphs so the compiler has opportunities to optimize both. For inference where gradients are not required, there are no training backward graphs to generate.
03:27 At this point, we have a representation of the computation that is ready for the backend compiler. This third stage is inductor, the default PyTorch backend compiler. Inductor takes the prepared graph and generates code for the hardware target. It can fuse compatible operations, choose how their work is scheduled and reduce unnecessary intermediate memory traffic.
03:53 For example, suppose a model adds a bias and then applies an activation function. Instead of launching separate kernels and storing an intermediate result, the compiler may be able to fuse those operations into one kernel. On GPUs, inductor commonly generates Triton kernels. On CPUs, it generates C++ code. It can also call optimized libraries such as those for matrix multiplication when that is the better implementation.
04:17 Not every operation is fused, though. Large matrix multiplications and other specialized operations may still use separate library kernels. The goal is to generate efficient code for the workload, not produce the smallest possible number of kernels. Now, what happens when Dynamo reaches code that it can't capture? With the default setting, it may create a graph break.
04:38 Dynamo finalizes the graph it has captured so far and sends it to the back end. It runs the unsupported code in regular Python and then resumes tracing. For example, a Python condition that depends on the value inside a tensor can cause a graph break. A call to item can also cause one depending on the code and compiler settings, the program may still run correctly, but the compiler now has smaller regions to optimize.
05:00 That can reduce fusion opportunities and add overhead at the boundaries. It does not mean that the entire model loses its speed up. There's also a stricter option setting full graph to true which requires the function to be captured as one graph and raises an error when that's not possible. So when performance is disappointing, graph breaks are one of the first things worth checking.
05:20 There is also a cost to the compilation itself. The first call pays the cost while later calls can reuse the same cache compiled code if the guards still match. If the input shapes or other assumptions change, PyTorch may need to recompile. Dynamic shape support can reduce some of these recompilations, but it does not eliminate every possible specialization.
05:40 That is why performance measurements should separate compilation time from steadystate execution time for GPU measurements. You also need proper synchronization so that you measure the work that actually complete it. So behind torch compile, remember the three stages. Dynamo captures the computation and records the conditions under which it can be reused.
06:01 AOT Autograd prepares the graph and handles the backward graph capture for training. Inductor generates code for the hardware. When compilation doesn't help, I would ask four questions. Did we capture enough of the model? Are graph breaks limiting optimizations? Are changing inputs causing recompilation? And are we measuring the compilation cost or the steadystate execution time?
06:24 Those questions will tell you much more than simply knowing that the model was compiled. The real value of torch compile is that it gives PyTorch the chance to optimize across operations that would otherwise be executed separately. Whether that opportunity turns into a speed up depends on the model and the workload.