Seonglae ChoSpeed up using JIT Compilation
Unlike torch.jit, this is a completely different JIT that performs kernel-level optimization by capturing the computational graph and passing it to the backend for Python execution
One issue is that when using forward hooks to extract intermediate hidden states, compile may skip or break hooks while optimizing the graph
Backends
- “inductor” default
- “nvfuser”
- "aot_eager” Pytorch Eager Mode but extract graph
- XLA for TPU
import torch model = MyModel() compiled_model = torch.compile(model, dynamic=True, fullgraph=False)
To avoid crashes, also set
fullgraph=False so PyTorch compiles only the safe parts. If your model takes inputs of varying sizes e.g., text sequences or images, use dynamic=True to ensure compatibility.