PyTorch normally runs tensor operations as Python reaches them: this is eager execution. For a model that needs faster execution, `torch.compile()` offers a small code change: wrap the model, then call the compiled model as usual. Introduced in PyTorch 2.0, it uses just-in-time compilation to optimize execution. section 17.1.5 section 1.5
Compile a model for inference
Load a model, move it to the chosen device, then wrap it: model = AutoModelForCausalLM.from_pretrained(checkpoint) model = model.to(device) model = torch.compile(model) The same model variable now holds the compiled model. Speedups depend on the model and hardware; compilation can occasionally slow execution down. section 17.1.5
Two optimizations help explain why compilation may be faster:
- Kernel fusion combines operations so they need fewer separate kernel calls. section 17.1.5
- Operator reordering can improve memory access and reduce intermediate memory use. section 17.1.5
Compilation also differs from `torch.export`. Export expects a graph it can capture; a branch chosen from an input value can make export fail. `torch.compile()` can instead break execution into smaller graphs around operations it cannot compile. Those graph breaks allow dynamic behavior, though they may cost performance. section 17.3 section 17.3.1
Listing 17.10 A model with dynamic conditional control flow
import torch
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class ConditionalModel(nn.Module):
def __init__(self):
super(ConditionalModel, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(10, 3)
def forward(self, x):
if x.sum() > 0:
x = self.fc1(x)
else:
x = self.fc2(x)
x = self.relu(x)
return x
model = ConditionalModel()
model = model.to(device)
model.eval()
The model chooses between two layers based on `x.sum()`. This input-dependent branch makes the difference between export and compilation concrete. section 17.3.1
