Python / PyTorch Fundamentals Interview Questions
What is torch.compile() and how does it speed up PyTorch model execution?
Introduced in PyTorch 2.0, torch.compile() performs just-in-time compilation of a model. Instead of executing each tensor operation eagerly (PyTorch's default), it captures the computation graph, fuses operations, and generates optimised kernels — primarily reducing GPU memory round-trips.
import torch import torch.nn as nn import time model = nn.Sequential( nn.Linear(1024, 1024), nn.GELU(), nn.Linear(1024, 512), nn.GELU(), nn.Linear(512, 10), ).cuda() # Compile the model â wraps it, does NOT change the API compiled_model = torch.compile(model) x = torch.randn(256, 1024).cuda() # First call triggers compilation (slow â may take 10-60 seconds) out = compiled_model(x) # Subsequent calls use the compiled, optimised kernels (fast) for _ in range(5): out = compiled_model(x) # Compilation modes â trade compile time for runtime speed model_default = torch.compile(model) # balanced model_reduce = torch.compile(model, mode="reduce-overhead") # less Python overhead model_max = torch.compile(model, mode="max-autotune") # slowest compile, fastest run # Benchmark comparison def benchmark(fn, x, n=100): for _ in range(5): fn(x) # warmup torch.cuda.synchronize() start = time.time() for _ in range(n): fn(x) torch.cuda.synchronize() return time.time() - start eager_time = benchmark(model, x) compiled_time = benchmark(compiled_model, x) print(f"Eager: {eager_time:.3f}s, Compiled: {compiled_time:.3f}s")
More Related questions...