Compile

PyTorch Compilation

torch.compile (PyTorch 2.0+)
JIT compiles Python/PyTorch code into optimized kernels for significant speedups.

Basic Usage

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 |

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

# 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

# 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)

# Debug compilation
model = torch.compile(model, fullgraph=False)  # Allow graph breaks

For Inference Optimization

# Combine with other optimizations
model = model.half()  # FP16
model = torch.compile(model, mode="reduce-overhead")
model.eval()

with torch.no_grad():
    output = model(input)