Eager Mode · Computation Graphs · Kernel Fusion
Learn through animations and interactive demos.
In eager mode, Python is the "driver". Every PyTorch operation executes immediately when Python hits that line — one by one, with Python overhead on every single GPU call.
Each op fires on the GPU the moment Python reaches that line. No buffering, no planning ahead.
Python calls into CUDA/C++ for each op. That Python overhead happens on every single operation.
Print tensors mid-computation, add breakpoints, inspect shapes. Great for research and experimentation.
class TinyNet(nn.Module): def forward(self, x): # (1) Python fires CUDA matmul RIGHT NOW x = self.linear1(x) # (2) Python fires CUDA relu RIGHT NOW x = self.relu(x) # (3) Python fires CUDA matmul RIGHT NOW x = self.linear2(x) return x # Python -> CUDA -> Python -> CUDA -> Python -> CUDA ...
Hit Run on both panels. Eager fires ops one at a time. Compiled fuses and fires them together.
torch.compile(model) replaces the Python-per-op loop with a single compiled execution plan. Three specialized stages run before any GPU kernel fires.
# One line changes everything model = torch.compile(model) output = model(x) # first call: traces + compiles (warm-up) output = model(x) # all subsequent: runs optimized compiled graph # Modes: torch.compile(model, backend="inductor") # default torch.compile(model, mode="reduce-overhead") # low latency torch.compile(model, mode="max-autotune") # max throughput
The forward graph is a DAG (Directed Acyclic Graph) — data flows in one direction: input tensors through operations to the output. Click Animate to watch data propagate node by node.
The backward graph is how the model learns. Gradients flow upstream through the same operations in reverse — computing how much each weight contributed to the error.
optimizer.zero_grad() # clear old gradients logits = model(x) # FORWARD -- build computation graph loss = criterion(logits, y) # LOSS -- scalar error value loss.backward() # BACKWARD -- dL/dw for all weights optimizer.step() # UPDATE -- w = w - lr * dL/dw # With torch.compile: fwd + bwd compiled as one unit model = torch.compile(model) # backward graph compiled once, reused every step
GPUs are fast at math but memory-bandwidth limited. Every separate kernel reads from and writes back to slow global GPU memory (HBM). Fusion keeps data in fast on-chip registers across multiple ops — eliminating those expensive round-trips.
relu, gelu, scale, add bias fuse trivially. Multiple element-wise ops collapse into one kernel pass through data.
softmax, layer_norm, mean combine a pointwise pass with a reduction. Fusing avoids writing partial results to HBM.
Fuse QK matmul + softmax + AV matmul into one tiled kernel. ~10x less memory than naive attention implementation.
Before loss.backward() can run, PyTorch must know what the backward graph is. It builds this during the forward pass by attaching a grad_fn to every output tensor. This chain of grad_fn nodes is the "tape" — the backward graph in disguise.
requires_grad=True, every op on it records itself as a grad_fn node on the output. After loss.backward(), these nodes are walked in reverse to compute all gradients. The tape is discarded after backward to free memory.import torch x = torch.randn(4, 4, requires_grad=True) w = torch.randn(4, 4, requires_grad=True) y = x @ w # y.grad_fn = MmBackward0 z = torch.relu(y) # z.grad_fn = ReluBackward0 print(z.grad_fn) # ReluBackward0 print(z.grad_fn.next_functions) # [(MmBackward0, 0)] # This chain IS the backward graph! loss = z.sum() loss.backward() # walks chain in reverse, computes x.grad + w.grad # Tape discarded after this to free memory
A graph break happens when TorchDynamo encounters something it cannot trace. It stops the current graph, runs that part eagerly (Python overhead returns), then starts a new graph. Too many breaks kill your speedup.
torch._dynamo.explain(model)(x) to see exactly where breaks occur and why.exp = torch._dynamo.explain(model)(x) print(exp.graph_count) # 1 = perfect print(exp.break_reasons) # why each break happened # Strict: error if any break occurs model = torch.compile(model, fullgraph=True) # Handle dynamic shapes without recompilation model = torch.compile(model, dynamic=True)
Kernel fusion works because of how GPU memory is structured. Data in different memory levels has wildly different speeds. Click each level to understand why this makes fusion so powerful.
Not every model benefits. Use this decision tree to find the right answer for your use case.
Fixed shapes, production workload, model >10M params, training throughput matters, warmup budget available.
Debugging, research, tiny models, variable batch sizes every call, single-sample latency critical.
Cross-platform export needed (ONNX/TensorRT), mobile deployment (torch.export), quantization pipeline required.
The profiler shows exactly which GPU kernels are running and how long they take. Eager = many narrow bars. Compiled = few wide fused bars. Click to animate the comparison.
from torch.profiler import profile, ProfilerActivity with profile(activities=[ProfilerActivity.CUDA]) as prof: model(x) print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10)) # Eager output: aten::mm 2.1ms | aten::relu 0.3ms | aten::add 0.4ms ... # Compiled: triton_fused_mm_relu_add 1.1ms <-- all ONE kernel! prof.export_chrome_trace("trace.json") # open chrome://tracing
torch.compile is not always faster. Here are real numbers to calibrate your expectations on A100 GPU, batch=32, FP16.
| Model | Eager (ms) | Compiled (ms) | Speedup | Notes |
|---|---|---|---|---|
| ResNet-50 | 8.2 | 3.1 | 2.6x | Element-wise ops fuse heavily |
| GPT-2 (124M) | 14.7 | 5.8 | 2.5x | Attention + FFN fusion |
| BERT-base | 11.3 | 4.6 | 2.5x | LayerNorm + GeLU fusion key |
| LLaMA-7B (fwd) | 182 | 74 | 2.5x | Flash Attention style |
| Simple 2-layer MLP | 1.1 | 1.4 | 0.8x | Compile overhead > savings |
| Variable batch sizes | 8.2 | 9.1 | 0.9x | Recompilation every step |
| Approach | Best For | Pros | Cons |
|---|---|---|---|
| Eager | Research, debugging | Full Python, easy debugging | Slow, no cross-op optimization |
| torch.compile | Production training + inference | 2-4x faster, no code changes | Warm-up cost, dynamic shape issues |
| torch.jit.script | Deployment, C++ runtime | Serializable, no Python needed | Strict Python subset, manual annotation |
| torch.export + ONNX | Cross-framework (TensorRT) | Runs anywhere | Complex export, limited ops |
4 questions to test what you have learned.
Everything at a glance.
| Concept | What It Is | Analogy | Impact |
|---|---|---|---|
| Eager Mode | Python calls CUDA one op at a time, immediately. | Reading a recipe one line at a time, then cooking | Baseline |
| torch.compile | Traces model into a graph, compiles via Inductor. | Read whole recipe, then cook optimally | 2-4x faster |
| TorchDynamo | Intercepts Python bytecode, captures FX Graph. | Stenographer recording what you cook, not how | Enables compilation |
| AOT Autograd | Derives fwd + bwd graphs before training begins. | Planning recipe AND cleanup before you start | Bwd 2-3x faster |
| Forward Graph | DAG of ops: input through operations to prediction. | Flowchart from ingredients to finished dish | Compiled once |
| Backward Graph | Reverse graph: chain rule computes dL/dw for all params. | Finding which ingredient made the dish wrong | Compiled once |
| Kernel Fusion | Multiple GPU ops merged into one kernel, data in registers. | Chopping + cooking in one pan, no switching | Massive BW saving |
| TorchInductor | Backend compiler generating fused Triton/CUDA code. | Master chef writing the most efficient method | Enables fusion |