PyTorch

Part I: The mental model

Python API            torch/                    nn.Module, optim, tensor methods
      |
      |  generated bindings + argument parsing
      v
Dispatcher            c10/ + ATen/core/dispatch/  per-op table, keyed by DispatchKeySet
      |
      |  key priority: Autocast -> Autograd -> backend
      v
ATen operators        aten/src/ATen/native/      native_functions.yaml schemas
      |
      |  structured kernels, TensorIterator
      v
Device kernels        native/cpu/ (SIMD loops), native/cuda/ (kernel launches)

  side channel written on the way down, replayed later:
Autograd tape         torch/csrc/autograd/       graph of Nodes, engine.cpp walks it
  optional layer that captures instead of interpreting:
Compiler              torch/_dynamo, _functorch, _inductor

The one-sentence identity: PyTorch is a runtime-recorded autograd tape sitting on an extensible operator dispatcher, with a compiler bolted on top that captures graphs opportunistically instead of demanding them up front. Every call like torch.add descends through the same stack: a thin Python surface, a dispatcher that picks the right kernel based on keys the tensors carry, an operator library (ATen) whose inventory is declared in one YAML file, and device kernels at the bottom. Autograd is not a separate mode; it is one of the dispatch keys, and its kernel records a graph node before re-dispatching to the layer below.

That stack explains the repository's shape. Pure Python lives under torch/; the C++ operator library is ATen under aten/src/ATen/; the small core both depend on (TensorImpl, dispatch keys, device abstractions) is c10/; and the C++ that binds it all to Python, autograd engine included, is torch/csrc/. Two YAML files are the skeleton key: native_functions.yaml declares roughly 2,600 operator schemas and which kernel implements each per backend, and derivatives.yaml declares roughly 700 derivative formulas (counts from main, July 2026). Code generation turns both into C++ at build time, which is why some classes you will meet, like AddBackward0, exist in no checked-in source file.

Hold on to the two directions of travel. The forward direction is a call descending the stack and, as a side effect, appending nodes to a graph. The backward direction is an engine walking that graph in reverse, where every derivative computation is itself an operator call that descends the same stack again. Once you see that symmetry, the rest of this chapter is detail.

Part II: Using it

Installation is one line on both Linux and macOS, and the selector on pytorch.org gives the exact command for a specific CUDA or ROCm build:

pip install torch          # CPU on macOS (plus MPS on Apple silicon), CUDA-enabled on Linux

The latest stable release as I write is 2.13.0 (July 2026), which supports Python 3.10 through 3.14. A first session should be a REPL, because the object model teaches itself:

>>> import torch
>>> x = torch.tensor([2.0], requires_grad=True)
>>> y = x * x
>>> y
tensor([4.], grad_fn=<MulBackward0>)
>>> y.backward()
>>> x.grad
tensor([4.])

That grad_fn=<MulBackward0> is the whole story in miniature: running the multiply recorded a graph node, and backward() replayed it. The full training idiom fits in a screen. You define a module, run a forward pass, compute a loss, call backward to populate gradients, and let an optimizer apply them:

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

x = torch.randn(64, 784)
y = torch.randint(0, 10, (64,))

logits = model(x)
loss = loss_fn(logits, y)
loss.backward()
opt.step()
opt.zero_grad()

Now the mistakes every beginner makes, each of which is a direct consequence of internals covered later. First, gradients accumulate; they are added into .grad, not assigned. Forgetting to clear them silently mixes batches:

# wrong: gradients from every past step pile up in .grad
for xb, yb in loader:
    loss_fn(model(xb), yb).backward()
    opt.step()

# right: clear before (or after) each step
for xb, yb in loader:
    opt.zero_grad()
    loss_fn(model(xb), yb).backward()
    opt.step()

Second, the graph is freed as backward consumes it, so calling backward twice on the same forward pass fails with RuntimeError: Trying to backward through the graph a second time. The fix is almost never retain_graph=True; it is realizing you wanted one combined loss:

# wrong: second backward walks an already-freed graph
loss_a.backward()
loss_b.backward()   # if loss_b shares any forward computation with loss_a: RuntimeError

# right: sum losses, walk the graph once
(loss_a + loss_b).backward()

Third, in-place operations can invalidate tensors that autograd saved for the backward pass. The version-counter machinery (Part V) catches this at backward time, not at the mutation site, which confuses people because the error appears far from the bug:

x = torch.randn(3, requires_grad=True)
y = x.sigmoid()     # sigmoid saves its output to compute its derivative
y.mul_(2)           # mutates the saved tensor in place
y.sum().backward()  # RuntimeError: one of the variables needed for gradient
                    # computation has been modified by an inplace operation

Deeper usage is layered onto the same loop rather than replacing it: .to("cuda") for devices, torch.autocast for mixed precision, torch.no_grad() for inference (which skips tape recording entirely, saving both time and the memory of saved activations). When the built-in operator set is not enough, you can join the tape yourself with a custom autograd.Function, which is exactly the interface the generated operator nodes implement:

class Cube(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x ** 3

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        return 3 * x * x * grad_out

y = Cube.apply(torch.tensor([2.0], requires_grad=True))

And since the 2.0 release the same model can be handed to the compiler with one call, with the eager code kept working unchanged:

compiled = torch.compile(model)
logits = compiled(x)   # first call compiles (slow), later calls run the compiled code

Part III: When it is the right tool

PyTorch is the right default for research and training almost unconditionally: define-by-run means your model is ordinary Python you can step through in a debugger, the ecosystem (Hugging Face, torchvision, the CUDA kernel community) assumes it, and torch.compile recovers most of the performance that graph-first frameworks used to hold over it. It is also a fine substrate for production training at any scale, which is what nanoGPT demonstrates at the small end and torchtitan at the large end.

The honest cases for alternatives: JAX, when your work is shaped like composed function transforms (grad of vmap of pmap) over pure functions, when you target TPUs, or when you want XLA's whole-program compilation as the default rather than an opt-in. TensorFlow persists mostly where its deployment story (TFLite, TF Serving, established mobile pipelines) is already entrenched. And for serving LLMs specifically, a dedicated inference engine beats a hand-rolled PyTorch loop, because serving is a scheduling-and-memory problem before it is a math problem.

That last point is the architecture-shaped warning. Eager PyTorch pays Python overhead and a kernel launch per operator, holds the GIL between launches, and does nothing to batch concurrent requests. Wrapping model(x) in a web handler works in a demo and collapses under load:

dangerous: request thread -> Python handler -> eager model(x) -> GPU
           (one request at a time, GIL contention, per-op launch overhead,
            GPU idle between tiny kernels)

safe:      requests -> queue -> batching engine (compiled/exported model,
           or a serving system like vLLM) -> GPU at high occupancy

The same warning generalizes: any hot loop that crosses the Python-operator boundary millions of times per second is fighting the framework's design. The escape hatches, in order of effort, are torch.compile, CUDA graphs (mode="reduce-overhead"), and export to a runtime.

Part IV: The full life of one loss.backward()

This is the core of the chapter. The specimen is the small net from Part II: Linear(784, 128) → ReLU → Linear(128, 10) into a cross-entropy loss, on CPU for concreteness. Before following the backward call, it pays to follow one forward operator all the way down, because backward will reuse every stage of that path.

Prologue: what one torch.add call traverses

torch.add(a, b) does not hit a Python implementation. The binding is C++ generated at build time by tools/autograd/gen_python_functions.py, and its first job is argument parsing (torch/csrc/utils/python_arg_parser.cpp): matching your Python objects against the operator's schemas and converting them to C++ types. From there the call enters the dispatcher (aten/src/ATen/core/dispatch/Dispatcher.h). Every tensor carries a DispatchKeySet, a 64-bit bitset in its TensorImpl; the dispatcher unions the keysets of all arguments and jumps to the kernel registered for the highest priority key present in the operator's table.

For a CPU float tensor with requires_grad=True, the winning key is AutogradCPU. That lands in a generated function (from tools/autograd/gen_variable_type.py plus the add entry in derivatives.yaml) which does three things: constructs an AddBackward0 node, wires its edges to the grad_fns of the inputs, and re-dispatches with the autograd keys masked off. The second dispatch resolves to CPU, which for add is a structured kernel: the schema in native_functions.yaml says structured_delegate: add.out, and add.out builds a TensorIterator (aten/src/ATen/TensorIterator.cpp) that handles broadcasting, type promotion, and memory-overlap checks, then runs the actual arithmetic from the one-line loop body in aten/src/ATen/native/ufunc/add.h, vectorized with SIMD on CPU or instantiated as a CUDA kernel by the ufunc codegen. One user-visible call, four layers, and the only handwritten math is one line.

Stage 1: the forward pass wrote the tape

By the time you hold loss, every operator in the forward pass has done what add did above. nn.Linear called F.linear, which is a CompositeImplicitAutograd operator: it has no derivative formula of its own and instead decomposes into addmm (bias-plus-matmul), whose derivative is declared. So the graph hanging off loss looks like this, which you can verify by printing loss.grad_fn and walking .next_functions:

loss  grad_fn = NllLossBackward0
                     |
              LogSoftmaxBackward0          (CrossEntropyLoss = log_softmax + nll_loss)
                     |
              AddmmBackward0               (second Linear)
              /      |       \
   AccumulateGrad  ReluBackward0  AccumulateGrad
   (bias2 leaf)        |          (weight2 leaf, via a transpose view)
                  AddmmBackward0  (first Linear)
                  /      |     \
        AccumulateGrad  (input x: no edge, requires_grad=False)  AccumulateGrad
        (bias1 leaf)                                             (weight1 leaf)

Each node is a C++ torch::autograd::Node (torch/csrc/autograd/function.h) holding edges to its producers and SavedVariables for whatever the derivative formula needs (for addmm, the inputs; for relu, the output). Leaf tensors, the parameters, have no grad_fn; instead each owns an AccumulateGrad node that graph edges point at. Note what was not built: no graph for x, because it does not require grad, and nothing was planned ahead of time. The tape only records what actually ran, which is why Python control flow differentiates for free.

Stage 2: from Python into the engine

loss.backward() is Tensor.backward in torch/_tensor.py, which forwards to torch.autograd.backward in torch/autograd/__init__.py. That function validates that loss is scalar (or that you passed an explicit gradient), constructs the seed gradient, a ones tensor shaped like loss, and calls _engine_run_backward (torch/autograd/graph.py), which crosses into C++ through Variable._execution_engine.run_backward in torch/csrc/autograd/python_engine.cpp. The GIL is released here; from this point the traversal is pure C++ unless a Python hook or a custom Function forces a round-trip back up.

Stage 3: Engine::execute builds the GraphTask

Engine::execute in torch/csrc/autograd/engine.cpp wraps the roots in a GraphRoot node and builds a GraphTask, the per-invocation state: an error flag, the outstanding-task count, and crucially a dependency map. A traversal from the roots counts, for every reachable node, how many incoming gradient contributions it must wait for. This is topological bookkeeping: ReluBackward0 must not fire until the AddmmBackward0 above it has produced its gradient, and a node feeding two consumers must wait for both. The engine also owns one worker thread and one ReadyQueue per device, spun up lazily; CPU work runs on the calling thread. For a GPU model, backward for each device runs on that device's dedicated engine thread.

Stage 4: the ready queue drains

Execution is a worklist algorithm. A NodeTask pairs a node with an InputBuffer, the accumulated gradients arriving at that node. The loop in Engine::thread_main pops a task and calls evaluate_function, which runs pre-hooks, calls the node's apply() with the buffered gradients, then distributes the outputs along the node's edges: for each edge, decrement the consumer's dependency count, add the gradient into the consumer's InputBuffer (this addition is where gradients from forked paths merge), and push the consumer onto the right device's ready queue once its count hits zero. For our net the order is forced: NllLossBackward0, then LogSoftmaxBackward0, then the second AddmmBackward0, then ReluBackward0, then the first AddmmBackward0, with AccumulateGrad nodes firing as their single dependencies resolve.

Stage 5: each backward op descends the stack again

What does AddmmBackward0::apply actually do? It is generated code (in the build-time file torch/csrc/autograd/generated/Functions.cpp) implementing the formulas from derivatives.yaml. For a matmul the real entry reads:

# tools/autograd/derivatives.yaml (main, July 2026)
- name: mm(Tensor self, Tensor mat2) -> Tensor
  self: mm_mat1_backward(grad, mat2, self.sym_sizes(), self.sym_strides(), self.layout(), 1)
  mat2: mm_mat2_backward(grad, self, mat2.sym_sizes(), mat2.sym_strides(), mat2.layout(), 1)

Those helpers bottom out in at::mm calls: grad-of-input is grad @ weight, grad-of-weight is input^T @ grad. Every one of these is a first-class ATen operator that goes through the dispatcher, exactly like the prologue's add, just with the autograd keys excluded so no tape is recorded about the backward pass itself (unless you asked for create_graph=True, in which case backward is taped too and second derivatives work). Before using any saved tensor, the SavedVariable unpacking checks its version counter, which is where the in-place error from Part II is actually raised. On CPU these ops run to completion; on CUDA they are asynchronous kernel launches on a stream, and the engine handles cross-device edges with stream synchronization.

Stage 6: AccumulateGrad and the return to Python

The traversal terminates at the four AccumulateGrad nodes (torch/csrc/autograd/functions/accumulate_grad.cpp). Each takes the arriving gradient and either installs it as the parameter's .grad (stealing the buffer when it can, to avoid a copy) or adds it in place if .grad already exists. This accumulate-not-assign choice is deliberate: it is what makes gradient accumulation over micro-batches, and multiple backward calls from independent losses, work with no extra machinery, at the price of the zero_grad() ritual. As nodes complete, the graph's saved tensors are freed (the "buffers have already been freed" behavior), the outstanding-task count hits zero, Engine::execute returns, the GIL is retaken, and your next line of Python runs. The optimizer then reads .grad in opt.step(), closing the loop. Total machinery engaged for our toy net: one graph task, one thread, about a dozen nodes, and a few dozen dispatched operator calls.

Part V: Internals deep dives

Deep dive: the dispatcher

The dispatcher answers one question millions of times per second: given this operator and these arguments, which function runs? Its data structures are small. Each TensorImpl (c10/core/TensorImpl.h) carries a DispatchKeySet (c10/core/DispatchKeySet.h), a 64-bit bitset. Each operator has an OperatorEntry (aten/src/ATen/core/dispatch/OperatorEntry.h) with a table mapping keys to kernels. A call unions the argument keysets with thread-local included/excluded sets, finds the highest-priority set bit that the operator has a kernel for, and jumps. Priority is the load-bearing part, because keys are layers, not alternatives:

highest priority
   Autocast*        cast inputs, re-dispatch (above autograd so casts get taped)
   Autograd*        record a Node, re-dispatch with autograd masked off
   ADInplaceOrView  version-counter bumps and view bookkeeping
   backend keys     CPU, CUDA, MPS, XPU, Sparse*, Quantized*, Meta, ...
lowest priority
(simplified: functorch, Python-mode, and tracing keys also slot into this order)

A typical call therefore passes through several kernels, each doing its cross-cutting work and re-dispatching below itself, like middleware. Think of the dispatcher as an open extensibility point rather than a switch statement: autograd, autocast, functorch transforms, __torch_dispatch__ modes, and out-of-tree backends are all just keys, which is how they compose without knowing about each other. An ordering subtlety worth memorizing: Autocast sits above Autograd precisely so the dtype casts it inserts are recorded on the tape and differentiated correctly; the comment in c10/core/DispatchKey.h says so explicitly.

The operator inventory is declarative. Here is the real schema for add, abbreviated from aten/src/ATen/native/native_functions.yaml:

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  structured_delegate: add.out
  variants: function, method
  dispatch:
    SparseCPU, SparseCUDA, ...: add_sparse
    MkldnnCPU: mkldnn_add
    ZeroTensor: add_zerotensor
  tags: [core, pointwise]

structured_delegate says the dense implementation lives in the add.out entry, a structured kernel where shape checking and output allocation are written once and shared between the functional, in-place, and out= variants. torchgen/ reads this file and generates the registration glue, the Python bindings, and for pointwise ufuncs even the CUDA kernels. Two practical consequences. First, "where is the code for operator X" always has the same answer: find its YAML entry; the entry names the C++ functions to read. Second, extending PyTorch does not mean forking it: a TORCH_LIBRARY_IMPL block can register kernels for an existing op under a new backend key, which is exactly how out-of-tree devices ship. The famous misconception to correct: there is no Python-level if/else choosing CPU versus GPU, and torch.add has no Python body at all; by the time a device is chosen you are two layers below Python.

Deep dive: autograd

Three structures carry the whole system. A Node (torch/csrc/autograd/function.h) is a differentiable operation's backward half: an apply() taking incoming gradients, plus next_edges pointing at producer nodes. An Edge is a (node, input index) pair, the index mattering for multi-input nodes. And each tensor that participates has AutogradMeta (torch/csrc/autograd/variable.h) holding either a grad_fn (non-leaf: the node that produced it) or a lazily created grad_accumulator (leaf). This is why .grad on a non-leaf is None by default: intermediate gradients flow through InputBuffers and are discarded unless you call retain_grad() or install a hook.

Saved tensors are the memory story. A derivative formula declares what it needs, and the generated node stores those as SavedVariables, which is why training uses so much more memory than inference: the activations live until backward consumes them. Each save snapshots the tensor's version counter, and unpacking rechecks it, turning silent in-place corruption into a loud error. Two corollaries people miss: activation checkpointing is just choosing not to save (recompute in backward instead), and torch.no_grad() is just excluding the autograd keys from dispatch, so nothing is recorded and nothing is saved.

The derivative formulas live in tools/autograd/derivatives.yaml, one entry per differentiable op, in a C++-ish DSL:

- name: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  self: handle_r_to_c(self.scalar_type(), grad)
  other: handle_r_to_c(other.scalar_type(), maybe_multiply(grad, alpha.conj()))

- name: relu(Tensor self) -> Tensor
  self: threshold_backward(grad, result, 0)

gen_variable_type.py turns each entry into a Node subclass and an autograd kernel. The engine that replays them (engine.cpp) is worth a careful read with the Stage 3 and 4 vocabulary in hand: GraphTask, dependency counting, ReadyQueue per device, NodeTask, InputBuffer, and the reentrancy machinery that handles backward-inside-backward. Misconceptions to correct explicitly: the tape is rebuilt every iteration, so autograd never "optimizes your graph"; it faithfully replays what ran, and anything smarter is the compiler's job. And define-by-run does not mean derivatives are computed by tracing numerical perturbations; every formula is symbolic, written once in YAML, and exact.

Deep dive: torch.compile, at altitude

Eager execution pays a price: Python between every kernel launch, and no optimization across operator boundaries. torch.compile attacks this without giving up define-by-run, via three named components, the design laid out in the PyTorch 2 paper at ASPLOS 2024. I describe them at concept level, verified against main in July 2026, because this corner of the repo moves fastest.

TorchDynamo (torch/_dynamo/, with the CPython hook in torch/csrc/dynamo/eval_frame.c) intercepts Python frame execution and symbolically walks your function's bytecode (symbolic_convert.py), extracting tensor operations into an FX graph and compiling the rest of the frame into resume functions. For every assumption it bakes in, it installs a guard (guards.py): a cheap runtime check on shapes, dtypes, closure values, module identity. On later calls, if the guards pass, the compiled artifact runs; if not, it recompiles or falls back. When Dynamo meets Python it cannot capture (a data-dependent branch on a tensor value, an unsupported builtin) it inserts a graph break: eager runs that piece, and compilation resumes after. AOTAutograd (torch/_functorch/aot_autograd.py) then traces through the autograd machinery ahead of time so the backward pass becomes a graph too, and functionalizes mutations away. TorchInductor (torch/_inductor/) lowers the graphs into its own IR (lowering.py), fuses operations (scheduler.py), and emits code: Triton kernels on GPU (codegen/triton.py), C++ with OpenMP on CPU. Chains of pointwise ops become single kernels, so intermediates never touch memory; this is the same memory-traffic logic that motivates FlashAttention, applied automatically.

The design value running through all three is graceful partiality: unsupported code degrades to eager instead of failing, which is what makes the one-line adoption story honest. The traps: the first call is slow by design (compilation happens then, and again per new dynamic-shape bucket or guard failure), so benchmark steady state; frequent recompiles usually mean a guard on something that varies per call; and a model riddled with graph breaks gets little fusion benefit, which torch._dynamo.explain will show you precisely.

Part VI: Reading the repository

Do not read breadth-first. The staged plan below follows the same descent as Part IV; every file named here exists on main as of July 2026.

Stage 0, the Python surface. Read torch/nn/modules/module.py (parameter registration, hooks, _call_impl), then torch/nn/modules/linear.py, then skim torch/_tensor.py for how Tensor.backward and friends delegate. Questions you should be able to answer: where do parameters actually live in a module, what does model(x) do before your forward runs, and why is Tensor.backward three lines long?

Stage 1, one operator in ATen. Find add.Tensor in aten/src/ATen/native/native_functions.yaml, read the surrounding schema conventions, then aten/src/ATen/native/BinaryOps.cpp, aten/src/ATen/native/ufunc/add.h, and aten/src/ATen/TensorIterator.h. Questions: what does structured_delegate buy, where does broadcasting happen, and why do most pointwise ops contain no loop code of their own?

Stage 2, the dispatch core. Read c10/core/DispatchKey.h top to bottom (the comments are a book chapter in themselves), then c10/core/DispatchKeySet.h, aten/src/ATen/core/dispatch/Dispatcher.h, and OperatorEntry.h. Questions: how is per-backend autograd (AutogradCPU vs AutogradCUDA) encoded in 64 bits, what is a fallback kernel, and what happens when no kernel is registered for the computed key?

Stage 3, autograd. Read tools/autograd/derivatives.yaml (skim, then find mm and relu), torch/csrc/autograd/variable.h, function.h, functions/accumulate_grad.cpp, and then engine.cpp, which is long but rewards patience. Questions: what exactly is saved for backward and where, how do two gradient paths into one node merge, and on which thread does backward for a CUDA tensor run?

Stage 4, codegen. Skim torchgen/model.py and torchgen/gen.py, plus tools/autograd/gen_variable_type.py and gen_python_functions.py. Question: for one operator, name every artifact generated from its YAML entries.

Stage 5, the compiler. torch/_dynamo/convert_frame.py and symbolic_convert.py for capture, guards.py, then torch/_functorch/aot_autograd.py, then torch/_inductor/compile_fx.py, lowering.py, scheduler.py. Questions: what is a guard concretely, why does AOTAutograd exist at all given eager autograd works, and what decides whether two nodes fuse?

Where not to start: torch/csrc/jit/, the legacy TorchScript stack, is large, largely superseded by torch.compile, and will teach you history rather than the present; the build system (Bazel/CMake files, caffe2 remnants) is a tar pit; and the generated code directories only exist after a source build, so do not go hunting for AddBackward0 in the checked-in tree.

Part VII: Hands-on labs

Each lab is runnable with a stock pip install torch on CPU unless noted, and each teaches one concept from the deep dives. Exact object addresses and log formats vary by version.

Lab 1: walk the tape by hand. Concept: the graph of Stage 1.

import torch, torch.nn as nn
model = nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 2))
loss = nn.functional.cross_entropy(model(torch.randn(5, 4)), torch.randint(0, 2, (5,)))
node = loss.grad_fn
while node:
    print(type(node).__name__, "->", [type(f[0]).__name__ for f in node.next_functions if f[0]])
    # descend the first edge that leads somewhere (skips leaf AccumulateGrad nodes)
    node = next((f[0] for f in node.next_functions if f[0] and f[0].next_functions), None)

Observe the chain NllLossBackward0, LogSoftmaxBackward0, AddmmBackward0, ReluBackward0, AddmmBackward0, with AccumulateGrad and transpose-view nodes hanging off the addmm nodes. Match each to the diagram in Part IV.

Lab 2: prove gradients accumulate. Concept: AccumulateGrad.

w = torch.ones(2, requires_grad=True)
(w * 3).sum().backward()
print(w.grad)        # tensor([3., 3.])
(w * 3).sum().backward()
print(w.grad)        # tensor([6., 6.])  <- added, not replaced
w.grad = None        # what zero_grad(set_to_none=True) does
(w * 3).sum().backward()
print(w.grad)        # tensor([3., 3.])

Lab 3: watch the dispatcher work. Concept: every forward and backward op is a dispatched ATen call.

from torch.utils._python_dispatch import TorchDispatchMode
import torch

class Log(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        print(func)
        return func(*args, **(kwargs or {}))

a = torch.randn(8, 8, requires_grad=True)
w = torch.randn(8, 8, requires_grad=True)
with Log():
    (a @ w).relu().sum().backward()

Observe forward ops (aten.mm.default, aten.relu.default, aten.sum.default), then the seed (aten.ones_like.default), then the backward ops (aten.expand.default, aten.threshold_backward.default, two aten.mm.default calls with transposes). This is Stage 5 made visible: backward is just more operator calls. The exact list varies slightly across releases.

Lab 4: trigger the version counter. Concept: SavedVariable integrity.

x = torch.randn(3, requires_grad=True)
y = x.sigmoid()
y.mul_(2)
try:
    y.sum().backward()
except RuntimeError as e:
    print(e)   # "...modified by an inplace operation..." with a version mismatch

Then rerun with y = x.relu() replaced by an op that saves its input instead of its output, and observe which in-place edits are legal. Predicting the answer from derivatives.yaml is the exercise.

Lab 5: see what the compiler generates. Concept: Dynamo capture and Inductor fusion. GPU makes it more dramatic but CPU works.

import torch

def f(x):
    return (x.sin() + x.cos()).relu() * 2

x = torch.randn(1_000_000)
print(torch._dynamo.explain(f)(x))     # graph count, graph break count: expect 1 graph, 0 breaks
# Then, from a shell:
#   TORCH_LOGS="output_code" python thisfile.py
# and read the generated C++ (CPU) or Triton (GPU): one fused kernel
# for the whole pointwise chain instead of five.

Measure steady-state timing with and without torch.compile(f) (skip the first call; that is compilation). The speedup depends entirely on your hardware and the op mix, so measure rather than assume.

Lab 6: cause and observe a recompile. Concept: guards.

# TORCH_LOGS="recompiles" python thisfile.py
import torch

@torch.compile
def g(x, flag):
    return x * 2 if flag else x * 3

x = torch.randn(10)
g(x, True); g(x, True)    # second call: guards pass, no recompile
g(x, False)               # guard on `flag` fails -> logged recompile

Observe the log naming the exact guard that failed. Then vary the tensor's shape instead and watch dynamic-shape behavior: typically one recompile to a symbolic-shape version, then stability.

Part VIII: Questions and model answers

Understanding checks. Try to answer aloud before reading each model answer.

1. What is PyTorch, in one sentence a systems person would accept?

A runtime-recorded autograd tape on top of an extensible operator dispatcher, with an opportunistic graph compiler layered above and a code-generated operator library below.

2. What does loss.backward() actually execute?

It seeds the recorded graph's root with a ones gradient, then the C++ engine builds a GraphTask, counts dependencies, and drains per-device ready queues, calling each node's backward function and merging fan-in gradients in InputBuffers, until AccumulateGrad nodes write the leaves' .grad fields. Every derivative computation is itself a dispatched ATen operator call.

3. Why do gradients accumulate instead of overwrite?

Because AccumulateGrad adds into .grad, which makes micro-batch gradient accumulation and multiple independent losses work with zero extra machinery. The cost is that training loops must call zero_grad(), and forgetting it silently blends stale gradients into every step.

4. Why does calling backward twice fail, and when is retain_graph the wrong fix?

The engine frees saved tensors as it consumes the graph, so a second traversal finds them gone. retain_graph=True is correct only when you genuinely need two traversals of the same forward computation; if you simply have two losses, summing them and walking once is cheaper and clearer.

5. What is a dispatch key, concretely?

A bit in a 64-bit set carried by every TensorImpl and combined with thread-local state at call time. Each operator holds a table from keys to kernels, and the call jumps to the kernel for the highest-priority key present, with layered kernels like autocast and autograd doing their work and re-dispatching below themselves.

6. Why is Autocast a higher-priority key than Autograd?

So the dtype casts autocast inserts happen before recording, which puts the cast operations onto the tape and makes their gradients (and dtype transitions) correct in backward. If autograd ran first, the tape would describe a computation that never actually executed.

7. Where would you look for the implementation of an unfamiliar operator?

Its entry in aten/src/ATen/native/native_functions.yaml, which names the implementing C++ function per backend or the structured delegate; the derivative, if any, is in tools/autograd/derivatives.yaml. No entry under dispatch usually means it is composite and decomposes into other ops.

8. Why does F.linear have no derivative formula?

It is CompositeImplicitAutograd: it decomposes into addmm and matmul, which have formulas. Autograd differentiates composites by taping their constituent ops, so only primitive operators need entries in derivatives.yaml.

9. Why is a non-leaf tensor's .grad None after backward?

Intermediate gradients flow through the engine's InputBuffers and are discarded once consumed; only leaves have AccumulateGrad nodes that persist gradients. Call retain_grad() on the intermediate, or register a hook, if you need its gradient.

10. A user reports "modified by an inplace operation" at backward time. How do you debug it?

Some op saved a tensor for its derivative and an in-place op bumped that tensor's version afterward. Read the error's mention of which output and node, find the in-place op (often an innocuous relu_, add_, or an activation reused buffer), and either make it out-of-place or reorder it. torch.autograd.set_detect_anomaly(True) will point at the recording site.

11. What are Dynamo guards and why do they exist?

Cheap runtime predicates (on shapes, dtypes, constants, module identities) attached to each compiled graph, encoding the assumptions baked into the compilation. They let compiled code be reused safely from arbitrary Python call sites: guards pass, run fast code; guards fail, recompile or fall back to eager.

12. Why does torch.compile need AOTAutograd when eager autograd already exists?

Eager autograd produces the backward incrementally at runtime, which the compiler cannot optimize ahead of time. AOTAutograd traces the autograd machinery once, ahead of execution, yielding a backward graph that Inductor can fuse and schedule like the forward, so both directions get compiled.

13. When would you reach for JAX instead?

When the work is naturally composed function transforms over pure functions, when TPUs are the target, or when whole-program XLA compilation and explicit PRNG/state threading fit the problem. PyTorch's edge is debuggability, the ecosystem, and mutation-friendly imperative code; JAX's is transform composability and a compiler-first execution model.

14. Why is the first call to a compiled model slow, and what do you do about it?

Compilation happens on first execution per guard-set: Dynamo capture, AOTAutograd tracing, Inductor codegen, and a Triton or C++ compile. Warm up before measuring or serving, keep shapes stable or rely on dynamic shapes, and watch TORCH_LOGS="recompiles" for guard churn that re-pays the cost repeatedly.

15. What does torch.no_grad() actually change?

It sets thread-local dispatcher state excluding the autograd keys, so operator calls skip the recording kernels entirely: no nodes, no saved activations. That is why it saves memory as well as time, and why it is orthogonal to model.eval(), which only flips module behavior like dropout and batch norm.

16. Why does PyTorch generate code from YAML instead of writing operators directly?

One schema entry fans out into bindings, dispatch registrations, autograd node classes, and per-variant kernels; with roughly 2,600 operators, hand-writing that surface would drift instantly. Single-source-of-truth declarations plus codegen keep the four layers mechanically consistent, at the cost of build complexity and generated classes you cannot grep for in the source tree.

Part IX: Design lessons

Declare once, generate everything. The YAML-plus-torchgen pattern keeps thousands of operators consistent across Python bindings, dispatch tables, and autograd nodes. The same move appears in protobuf/IDL-driven RPC stacks, database system catalogs, and syscall table generation: wherever a wide surface must stay mechanically in sync, a declarative source of truth beats discipline.

An open dispatch table is an extensibility contract. Making backends and cross-cutting features keys in a prioritized table lets autograd, autocast, vmap, and out-of-tree hardware compose without mutual knowledge. This is the VFS layer in operating systems, middleware chains in web frameworks, and LLVM's pass infrastructure: the framework owns the traversal, plugins own the behavior.

Record-replay beats declare-ahead when developer experience matters. Define-by-run won because the tape costs a little runtime overhead but makes models ordinary, debuggable programs. Write-ahead logs, tracing JITs, and record-and-replay debuggers make the same trade: capture what actually happened, then exploit the recording.

Graceful partiality makes compilers adoptable. Dynamo's graph breaks and eager fallback mean the compiler handles what it can and never blocks what it cannot, so adoption is one line rather than a rewrite. JavaScript JITs (deoptimization to interpreter) and query optimizers with fallback plans embody the same principle: a partial fast path plus a total slow path.

Guard cheaply what you cannot prove statically. Instead of demanding static types and shapes, compiled PyTorch assumes, guards, and recompiles on violation. Speculative execution in CPUs, optimistic concurrency control in databases, and inline caches in dynamic-language VMs are all this pattern: bet on stability, verify cheaply, pay only on the rare miss.

Part X: Memorization framework

The one-sentence summary: PyTorch runs every operator call through a keyed dispatcher whose autograd layer records a graph that an engine later replays in reverse, and torch.compile captures both directions into fused kernels when it can.

forward:   torch.add -> arg parser -> dispatcher -> Autograd key (record Node)
           -> backend key -> TensorIterator -> SIMD/CUDA loop

backward:  loss.backward() -> Engine::execute -> GraphTask deps
           -> ReadyQueue -> Node::apply (formulas from derivatives.yaml)
           -> dispatched kernels -> AccumulateGrad -> .grad -> opt.step()

compile:   Dynamo (capture + guards) -> AOTAutograd (joint fwd/bwd graph)
           -> Inductor (fuse, codegen Triton/C++)

The chain mapped to source:

arg parsing        torch/csrc/utils/python_arg_parser.cpp
dispatcher         c10/core/DispatchKey(Set).h, ATen/core/dispatch/Dispatcher.h
operator schemas   aten/src/ATen/native/native_functions.yaml
kernels            aten/src/ATen/native/{cpu,cuda}/, TensorIterator.cpp
derivatives        tools/autograd/derivatives.yaml
engine             torch/csrc/autograd/engine.cpp, functions/accumulate_grad.cpp
compiler           torch/_dynamo/, torch/_functorch/aot_autograd.py, torch/_inductor/

Memorize these blocks:

• Four directories: torch/ Python, aten/ operators, c10/ core types and keys, torch/csrc/ the C++ bridge and autograd engine.
• Two YAML files: native_functions.yaml (what operators exist, ~2,600) and derivatives.yaml (how ~700 of them differentiate); codegen turns both into C++.
• Five engine nouns: GraphTask, ReadyQueue, NodeTask, InputBuffer, AccumulateGrad; dependency counts gate execution, InputBuffer adds merge fan-in, AccumulateGrad writes .grad.
• Key order that matters: Autocast above Autograd above ADInplaceOrView above backend keys, each re-dispatching below itself.
• Compile trio: Dynamo captures with guards and breaks, AOTAutograd makes backward a graph, Inductor fuses and emits Triton or C++.

Part XI: Papers and further reading

The ideas in this walkthrough trace back to a short list of papers, and each one rewards a direct read. Where this site develops the same idea in depth, the companion link points there.

1. Paszke et al., PyTorch, An Imperative Style, High-Performance Deep Learning Library, NeurIPS 2019. The framework paper, stating the define-by-run design values this chapter keeps returning to. The PyTorch in the wild page on this site shows the resulting idioms in working code.
2. Ansel et al., PyTorch 2, Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, ASPLOS 2024. The TorchDynamo, AOTAutograd, and TorchInductor design of Part V, written by the team that built it.
3. Paszke et al., Automatic Differentiation in PyTorch, NIPS 2017 Autodiff Workshop. The short early paper on the tape and the variable system, worth reading as the seed of Part IV.
4. Baydin et al., Automatic Differentiation in Machine Learning, a Survey, JMLR 2018. The clearest map of forward and reverse mode, and of where a runtime-recorded tape sits in that landscape.
5. Li et al., PyTorch Distributed, Experiences on Accelerating Data Parallel Training, VLDB 2020. Gradient bucketing and communication overlap in DDP, the first rung of the scaling ladder climbed in the parallel computing class.
6. Zhao et al., PyTorch FSDP, Experiences on Scaling Fully Sharded Data Parallel, 2023. How sharded data parallelism lives inside the eager runtime described here, and the backbone of the torchtitan walkthrough.
7. Tillet et al., Triton, an Intermediate Language and Compiler for Tiled Neural Network Computations, MAPL 2019. The kernel language Inductor emits on GPU, covered in the Triton walkthrough.
8. Micikevicius et al., Mixed Precision Training, ICLR 2018. The recipe behind torch.autocast, which the dispatcher runs as a key above autograd so the casts land on the tape.
9. Dao et al., FlashAttention, Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022. The memory-traffic argument that Inductor's pointwise fusion applies automatically, derived in the FlashAttention walkthrough.

Part XII: Final takeaway

If you internalize one page of this chapter, make it Part IV: one call descending four layers on the way in, one engine walking a recorded graph on the way back, and every backward step re-entering the same four layers. The ML implementations section builds the numerical pieces these layers compute, and nanoGPT is the ideal codebase to watch this machinery carry a real model.