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

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)