compile
**PyTorch Compilation**
**torch.compile (PyTorch 2.0+)**
JIT compiles Python/PyTorch code into optimized kernels for significant speedups.
**Basic Usage**
```python
import torch
model = YourModel()
model = torch.compile(model) # That's it!
# First run is slow (compilation)
# Subsequent runs are fast
output = model(input)
```
**Compilation Modes**
**Available Modes**
| Mode | Speedup | Compile Time | Use Case |
|------|---------|--------------|----------|
| default | Moderate | Moderate | General use |
| reduce-overhead | High | Higher | Low latency |
| max-autotune | Highest | Very high | Benchmarking |
```python
model = torch.compile(model, mode="reduce-overhead")
```
**How It Works**
1. **Trace**: Capture computation graph (torch.fx)
2. **Optimize**: Apply graph optimizations
3. **Codegen**: Generate optimized kernels (Triton)
4. **Cache**: Reuse compiled kernels
**Benefits**
- **Kernel fusion**: Combine multiple ops into one
- **Memory optimization**: Reduce intermediate tensors
- **Automatic**: No manual optimization needed
**Performance Example**
```python
# Before compile
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b")
# ~45 tokens/second
# After compile
model = torch.compile(model)
# ~60+ tokens/second (30% faster)
```
**Considerations**
**Compilation Overhead**
- First run includes compilation time
- For inference: warm up before benchmarking
- Compilation cached within process
**Dynamic Shapes**
```python
# Disable for dynamic shapes (variable-length sequences)
torch._dynamo.config.dynamic_shapes = True
# Or mark dynamic dimensions
model = torch.compile(model, dynamic=True)
```
**Compatibility**
Not all operations are supported. Check for:
- Custom CUDA kernels
- Some external libraries
- Graph breaks (fallback to eager mode)
```python
# Debug compilation
model = torch.compile(model, fullgraph=False) # Allow graph breaks
```
**For Inference Optimization**
```python
# Combine with other optimizations
model = model.half() # FP16
model = torch.compile(model, mode="reduce-overhead")
model.eval()
with torch.no_grad():
output = model(input)
```