1.7 Automatic Differentiation & PyTorch Internals¶

Backpropagation is the algorithm that makes deep learning tractable, but writing it by hand for every new architecture is error-prone, slow, and frankly miserable. Automatic differentiation (autodiff) is the engineering discipline that solves this: given any composition of differentiable operations expressed as code, autodiff computes exact gradients — not finite differences, not symbolic algebra output — mechanically and efficiently. PyTorch’s autograd engine is the most widely used reverse-mode autodiff system in research and production. Understanding it at the mechanism level pays dividends every time you write a custom loss, a non-standard layer, or debug a gradient that quietly went to zero.

This chapter covers the full stack: the mathematics of reverse-mode autodiff and the tape metaphor, the PyTorch computation graph and its memory model, leaf tensors and the .grad accumulation protocol, what the graph actually stores (saved tensors, and how to count their bytes), the hook protocol that DDP and FSDP are built on, no_grad and inference mode, custom autograd.Function, and the lower layers of PyTorch (the dispatcher, ATen, views vs. copies, contiguity, and broadcasting semantics). We connect the theory to the practice with a fully-worked custom function example and numerical traces you can follow by hand.

We assume you have read Neural Networks From Scratch: MLPs & Backprop and are comfortable with the chain rule. We also reference Calculus, Optimization & Convexity for the Jacobian formalism and GPU Architecture & The Memory Hierarchy for the hardware context that makes contiguity matter.

Why Autodiff Exists: The Three Alternatives and Their Failures¶

Before diving into the mechanism, it is worth naming the alternatives that autodiff replaced.

Manual backprop requires the programmer to derive and implement the gradient of every operation. This was the norm in the symbolic-layer era (Theano, early Caffe). It is correct when done carefully, but it couples forward and backward code, makes architectural experiments tedious, and is a perennial source of subtle bugs.

Numerical differentiation (finite differences) approximates the derivative as:

This requires \(2N\) forward passes for \(N\) parameters — completely impractical for networks with billions of parameters. It is still useful for gradient checking (verifying autodiff implementations), where you compare the autodiff gradient against a small finite-difference estimate for a handful of parameters.

Symbolic differentiation (as in computer algebra systems like Mathematica) manipulates expression trees algebraically. The output is a symbolic formula for the derivative. It produces exact answers but suffers from expression swell — derivatives of composite functions grow exponentially in size with depth, and the resulting code is typically far slower than an equivalent imperative implementation.

Reverse-mode autodiff threads the needle: it computes exact derivatives (to floating-point precision), scales in \(O(1)\) forward passes regardless of \(N\), and operates on ordinary imperative code. Its only real cost is memory for the intermediate activations stored during the forward pass — a cost we will see how to reduce with gradient checkpointing (covered in Memory-Efficient Training: Checkpointing, Offloading & LoRA Math).

Reverse-Mode Autodiff: The Tape¶

The Jacobian-Vector Product View¶

Let \(f : \mathbb{R}^n \to \mathbb{R}^m\) be a differentiable function. Its derivative at a point \(x\) is a linear map \(Df(x) : \mathbb{R}^n \to \mathbb{R}^m\), represented as the \(m \times n\) Jacobian matrix \(J\). For most neural network losses, \(m = 1\), so the Jacobian is a \(1 \times n\) row vector — the gradient.

Reverse-mode autodiff computes the vector-Jacobian product (VJP):

where \(\bar{y}\) is the upstream gradient (a \(1 \times m\) row vector) and \(\bar{x}\) is the resulting gradient with respect to \(x\). When \(m = 1\) and \(\bar{y} = 1\), this recovers the ordinary gradient \(\nabla_x f\).

The key insight is that we never materialize \(J\) itself — we only ever compute VJPs. This makes reverse-mode autodiff efficient when the output dimension \(m\) is small (e.g., \(m = 1\) for scalar losses) regardless of how large \(n\) is.

The Tape Metaphor¶

During the forward pass, autograd records every operation applied to tensors that require gradients, building a directed acyclic graph (DAG). Edges point from outputs to inputs (in the backward direction). Nodes are Function objects that know how to compute the VJP for their operation. This data structure is called the tape or Wengert list (after Robert Wengert, who described it in 1964).

During the backward pass, we traverse the tape in reverse topological order. At each node, we: 1. Receive the upstream gradient \(\bar{y}\) from the node above. 2. Call the node’s backward function to compute the VJP: \(\bar{x} = \bar{y}^\top J\). 3. Accumulate \(\bar{x}\) into the gradient of the input tensor and pass it downstream.

After loss.backward() returns, every leaf tensor with requires_grad=True has its .grad field populated with the accumulated gradient.

PyTorch Autograd: The Computation Graph¶

Tensors, requires_grad, and Leaf Nodes¶

Every PyTorch tensor has a requires_grad flag. When True, operations on it are tracked. A tensor is a leaf if it was created directly by user code (e.g., nn.Parameter, or a tensor created with requires_grad=True), rather than being the output of some tracked operation. After backward(), only leaf tensors’ .grad fields are populated; intermediate (non-leaf) tensors’ gradients are not retained by default (to save memory).

import torch

# Leaf tensor: created directly, not the result of an op
x = torch.tensor([2.0, 3.0], requires_grad=True)
w = torch.tensor([0.5, -1.0], requires_grad=True)

# Non-leaf (intermediate) tensor: result of an operation
y = (x * w).sum()   # y = 2*0.5 + 3*(-1.0) = 1.0 - 3.0 = -2.0

print(x.is_leaf)   # True
print(y.is_leaf)   # False
print(y.grad_fn)   # <SumBackward0 object at 0x...>

y.backward()

# Only leaf gradients are populated
print(x.grad)   # tensor([ 0.5, -1.0])  -- dy/dx_i = w_i
print(w.grad)   # tensor([ 2.0,  3.0])  -- dy/dw_i = x_i
# y.grad would be None (non-leaf, gradient not retained)

The grad_fn attribute is a reference to the Function node that created this tensor. Following .grad_fn.next_functions traverses the graph toward the leaves.

The grad_fn Graph¶

# Inspecting the graph manually
a = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(4.0, requires_grad=True)
c = a * b           # MulBackward0
d = c + a           # AddBackward0
e = d ** 2          # PowBackward0

print(e.grad_fn)                          # PowBackward0
print(e.grad_fn.next_functions)           # ((AddBackward0, 0),)
print(e.grad_fn.next_functions[0][0].next_functions)
# ((MulBackward0, 0), (AccumulateGrad, 0))  ← 'a' appears twice!

Notice that a appears twice in the graph (once as an input to c = a*b and once as the second input to d = c+a). PyTorch correctly accumulates both gradient contributions into a.grad.

Gradient Accumulation and .grad_fn Ownership¶

By default, .grad accumulates (adds) across multiple .backward() calls. This is intentional and exploited by gradient accumulation in training:

optimizer.zero_grad()          # clear accumulated grads
for mini_batch in accumulation_steps:
    loss = model(mini_batch) / accumulation_steps
    loss.backward()            # accumulates into .grad
optimizer.step()               # update once with the full-batch gradient

If you forget zero_grad(), gradients from the previous step corrupt the current one — a classic bug.

What the Graph Stores: Saved Tensors and Their Cost¶

The tape is not free. Each Function node keeps alive exactly the tensors its backward formula needs — its saved tensors — and those references are why activation memory, not parameter memory, dominates the peak footprint of a training step. PyTorch exposes them for inspection through underscore-prefixed attributes on grad_fn, which is the fastest way to answer “what is this op actually holding on to?”:

import torch

x = torch.randn(1024, 512, requires_grad=True)
w = torch.randn(512, 512, requires_grad=True)

y = x @ w                      # MmBackward0 needs BOTH operands: A_bar = C_bar B^T, B_bar = A^T C_bar
print(y.grad_fn._saved_self.shape)    # torch.Size([1024, 512])  -- x is pinned by the graph
print(y.grad_fn._saved_mat2.shape)    # torch.Size([512, 512])   -- w is pinned by the graph

r = torch.relu(y)              # ReluBackward0 saves only the OUTPUT (mask = output > 0)
print([a for a in dir(r.grad_fn) if a.startswith("_saved")])   # ['_saved_result']

e = y.exp()                    # ExpBackward0: d/dx e^x = e^x, so the output suffices
print(e.grad_fn._saved_result.shape)  # torch.Size([1024, 512])

Two design rules fall out of this. First, prefer backward formulas that read the output, as StableSigmoid below does: relu, exp, sigmoid, and tanh all save one tensor instead of two. Second, matmuls are the expensive nodes — every Linear in your model pins its input activation for the whole backward pass.

Because saved tensors are just Python-visible objects flowing through the engine, you can intercept them. torch.autograd.graph.saved_tensors_hooks(pack, unpack) installs a pair of callbacks: pack runs when a tensor is saved (return anything you like — a CPU copy, a quantized blob, a filename), and unpack runs when backward needs it back. This one mechanism is the substrate under CPU activation offloading and quantized-activation training:

from torch.autograd.graph import saved_tensors_hooks

# Offload every saved activation to host memory, fetch it back in backward.
# pack may return ANY object; unpack must turn it back into the tensor.
# (torch.autograd.graph.save_on_cpu(pin_memory=True) is the built-in version of this.)
def pack(t):
    return (t.device, t.to("cpu", non_blocking=True))

def unpack(packed):
    device, t = packed
    return t.to(device, non_blocking=True)

with saved_tensors_hooks(pack, unpack):
    out = model(batch)          # graph built here holds CPU copies, not GPU tensors
loss = out.sum()
loss.backward()                 # unpack runs as each node is reached

Trading PCIe bandwidth for HBM this way is only a win when the transfer overlaps compute; see Memory-Efficient Training: Checkpointing, Offloading & LoRA Math for when offloading beats recomputation.

Hooks: Intercepting Gradients As They Flow¶

The engine calls user callbacks at two points, and essentially every distributed-training library is built on them.

• tensor.register_hook(fn) fires when that tensor’s gradient has been computed. Return None to merely observe, or return a tensor to replace the gradient that continues downstream. This is how you implement per-tensor clipping, gradient reversal layers, or a probe that logs which layer first goes NaN.
• tensor.register_post_accumulate_grad_hook(fn) (PyTorch 2.1+) fires on a leaf after .grad has been updated — that is, once the parameter’s gradient is final for this backward. nn.Module.register_full_backward_hook gives the same idea at module granularity.

import torch

lin = torch.nn.Linear(4, 4)
x = torch.randn(2, 4)

# Observe (and optionally rewrite) the gradient as it reaches this tensor.
h1 = lin.weight.register_hook(lambda g: print("weight grad norm:", g.norm().item()))

# Fires once, after lin.weight.grad is final. DDP attaches its all-reduce here.
def on_grad_ready(param):
    print("grad ready for", tuple(param.shape))

h2 = lin.weight.register_post_accumulate_grad_hook(on_grad_ready)

lin(x).sum().backward()
h1.remove()                     # hooks are handles -- always remove them
h2.remove()

Why this matters beyond debugging: DistributedDataParallel registers exactly this kind of hook on every parameter, so that as soon as a bucket of parameters has its gradients ready, the all-reduce for that bucket launches while the rest of the backward is still running. That overlap of communication with computation is the entire reason DDP is not bandwidth-bound — see Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP, which reconstructs the bucketing engine from these primitives. FSDP uses the same hook points to trigger the reduce-scatter and to free resharded parameters. For the debugging use, a register_hook that checks torch.isfinite(g).all() per layer is the cheap always-on complement to the anomaly detector at the end of this chapter.

torch.no_grad, inference_mode, and Detach¶

torch.no_grad()¶

Inside a no_grad context, autograd does not record operations even for tensors with requires_grad=True. No grad_fn is attached to outputs, and no tape is built. Use this for inference and for optimizer steps (parameter updates must not themselves be differentiated).

with torch.no_grad():
    output = model(input)   # no graph built; saves memory and compute

torch.inference_mode()¶

Introduced in PyTorch 1.9, inference_mode is a stronger version of no_grad: tensors created inside it are marked as “inference tensors” and cannot be used in a future requires_grad computation. This allows PyTorch to skip additional bookkeeping. Prefer it over no_grad for pure inference paths.

with torch.inference_mode():
    logits = model(input_ids)
    probs = torch.softmax(logits, dim=-1)

.detach()¶

.detach() returns a new tensor that shares the same storage but is detached from the computation graph — it has requires_grad=False and no grad_fn. Common uses:

• Logging or visualization of activations without building a graph.
• Stopping gradient flow in architectures like target networks in RL (where we want the target to be fixed).
• The “stop-gradient” trick in self-supervised learning (BYOL, SimSiam).

# Stop-gradient for target network
with torch.no_grad():
    target_features = target_encoder(x)  # equivalent to detach here

# Or explicitly:
target_features = online_encoder(x).detach()

Custom autograd.Function¶

PyTorch’s built-in operations cover almost every need, but sometimes you need a custom forward/backward pair: a fused kernel, a numerically stable reformulation, or a straight-through estimator. torch.autograd.Function gives you a clean interface to plug into the autograd engine.

The Function Interface¶

import torch
from torch.autograd import Function

class MyOp(Function):
    @staticmethod
    def forward(ctx, *inputs):
        # ctx is a context object for stashing tensors for backward
        # return output tensor(s)
        ...

    @staticmethod
    def backward(ctx, *grad_outputs):
        # grad_outputs: upstream gradients (one per forward output)
        # return gradient tensors (one per forward input, or None if not differentiable)
        ...

The ctx object is the bridge: ctx.save_for_backward(...) stashes tensors (only tensors, not Python scalars), and ctx.saved_tensors retrieves them in backward. To stash non-tensor values, assign them as attributes (ctx.alpha = alpha).

Example: Numerically Stable Sigmoid with Custom Backward¶

The naive sigmoid \(\sigma(x) = 1/(1+e^{-x})\) can overflow for large negative \(x\) (exp of large positive becomes inf) or lose precision for large positive \(x\). The stable version clips large magnitudes and uses torch.sigmoid in practice, but here we implement it from scratch with a custom backward to illustrate the interface:

import torch
from torch.autograd import Function


class StableSigmoid(Function):
    """
    Sigmoid with a numerically stable forward and analytic backward.
    We store only the output (not input) for memory efficiency,
    since d_sigma/dx = sigma(x) * (1 - sigma(x)).
    """

    @staticmethod
    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
        # Use PyTorch's stable sigmoid internally
        y = torch.sigmoid(x)
        # Save output, not input -- saves memory for large activations
        ctx.save_for_backward(y)
        return y

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
        (y,) = ctx.saved_tensors
        # Jacobian of sigmoid: dy/dx = y * (1 - y)
        # VJP: grad_input = grad_output * dy/dx  (element-wise for elt-wise ops)
        grad_input = grad_output * y * (1.0 - y)
        return grad_input


# Register as a callable
stable_sigmoid = StableSigmoid.apply


# --- Test correctness against autograd ---
torch.manual_seed(0)
x = torch.randn(4, requires_grad=True)
x_ref = x.detach().clone().requires_grad_(True)

y = stable_sigmoid(x)
y_ref = torch.sigmoid(x_ref)

# Forward agreement
assert torch.allclose(y, y_ref, atol=1e-6), "Forward mismatch"

# Backward agreement via gradient checking
y.sum().backward()
y_ref.sum().backward()
assert torch.allclose(x.grad, x_ref.grad, atol=1e-6), "Backward mismatch"

print("x      :", x.detach().numpy().round(4))
print("sigma  :", y.detach().numpy().round(4))
print("grad   :", x.grad.numpy().round(4))
# x      : [ 1.5410 -0.2934 -2.1788  0.5684]
# sigma  : [0.8238 0.4271 0.1017 0.6387]
# grad   : [0.1449 0.2446 0.0912 0.2307]  -- sigma*(1-sigma)

Example: Straight-Through Estimator (STE)¶

The STE is a classic trick used in quantization-aware training (QAT) and binary neural networks. The forward pass applies a non-differentiable rounding operation; the backward pass pretends the function was the identity, passing the upstream gradient through unchanged:

class StraightThroughRound(Function):
    """
    Forward:  y = round(x)   (non-differentiable: gradient is 0 a.e.)
    Backward: dy/dx = 1      (identity straight-through estimator)
    """

    @staticmethod
    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
        return x.round()

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
        # Pass gradient through unchanged -- the "straight-through" trick
        return grad_output


ste_round = StraightThroughRound.apply

x = torch.tensor([0.3, 1.7, -0.6], requires_grad=True)
y = ste_round(x)
print(y)           # tensor([ 0.,  2., -1.])
y.sum().backward()
print(x.grad)      # tensor([1., 1., 1.])  -- as if dy/dx = 1 everywhere

gradcheck: Numerically Verifying Custom Backwards¶

PyTorch provides torch.autograd.gradcheck to compare your custom backward against finite differences:

from torch.autograd import gradcheck

# Use double precision for numerical stability of finite differences
x_check = torch.randn(3, dtype=torch.float64, requires_grad=True)
result = gradcheck(stable_sigmoid, (x_check,), eps=1e-6, atol=1e-4)
print(f"gradcheck passed: {result}")  # True

Always run gradcheck on new Function implementations — it catches sign errors, missing factors, and wrongly accumulated terms.

When autograd.Function Is Not Enough: torch.library¶

autograd.Function is the right tool for eager-mode research code, and it is how FlashAttention and most Triton kernels expose a differentiable Python surface. It has one important limitation in the 2026 stack: TorchDynamo cannot see through an arbitrary Function whose forward calls opaque code (a raw CUDA launch, a @triton.jit kernel), so dropping one into a torch.compiled model risks a graph break that costs you the surrounding fusions. The supported alternative is to register the operation as a real PyTorch operator with torch.library.custom_op (or torch.library.triton_op for Triton) and attach its gradient with torch.library.register_autograd — then the op is visible to torch.compile, torch.export, and the dispatcher like any ATen op. Writing GPU Kernels with Triton and CUDA Programming Essentials for ML Engineers walk through both registrations end to end.

Two smaller modernizations of the interface are worth knowing. PyTorch 2.x supports splitting forward(ctx, ...) into a ctx-free forward(...) plus a separate setup_context(ctx, inputs, output) staticmethod; this separation is what lets the same Function be transformed by torch.func (set generate_vmap_rule = True, or supply an explicit vmap staticmethod, to make it vmap-compatible). And ctx.needs_input_grad — a tuple of booleans, one per forward input — lets backward skip computing gradients nobody asked for, returning None instead; on a fused kernel that can halve the backward cost.

Worked Numerical Example: Forward + Backward Trace¶

import torch

x  = torch.tensor([[1.0], [0.0]])               # (2,1)
W1 = torch.tensor([[2., 1.], [-1., 3.]], requires_grad=True)
w2 = torch.tensor([[0.5], [0.5]], requires_grad=True)

z1 = W1 @ x
h  = torch.relu(z1)
z2 = (w2.T @ h).squeeze()
L  = 0.5 * z2 ** 2

L.backward()

print("W1.grad:\n", W1.grad)
# [[0.5, 0.],
#  [0. , 0.]]
print("w2.grad:\n", w2.grad)
# [[2.],
#  [0.]]

The PyTorch Dispatcher and ATen¶

Understanding what happens below autograd helps when you need to write custom C++ extensions, debug shape mismatches, or understand torch.compile’s transformations.

Layers of the PyTorch Stack¶

The dispatcher is a routing table. Every operation is registered under one or more “dispatch keys” (tags attached to tensors based on device/dtype/layout). When you call torch.mm(a, b), the dispatcher inspects the keys of a and b and calls the appropriate backend implementation. This architecture enables:

• Backend extensibility: XLA, MPS, custom hardware backends plug in without touching existing code.
• Transforms: torch.compile, vmap, grad (functorch) all work by inserting themselves at a dispatch key layer, intercepting operations.
• Operator overriding: You can register a custom kernel for a specific (op, backend) pair.

The crucial thing to internalize is that autograd is itself a dispatch key, not a special case. Keys are ordered, and a call falls through them from highest to lowest priority. Autograd sits above the backend keys: its kernel for mm allocates the MmBackward0 node, wires up next_functions, and then redispatches the same call to the next key down (CUDA, CPU, …) to actually compute the numbers. That is why no_grad is cheap — it simply excludes the Autograd key from the dispatch set, so the call lands straight on the backend kernel with no node allocated.

The same layering explains mixed precision. torch.autocast inserts an Autocast key above Autograd, whose kernel for each listed op casts the operands to bf16/fp16 before redispatching. Because the cast happens above autograd, the tensors the graph saves are the cast (bf16) ones — which is exactly why autocast reduces activation memory, and also why an unsafe op left on autocast’s fp32 list still saves fp32 activations. See Mixed Precision, bf16 & FP8 Training.

ATen and the Operator Schema¶

ATen is PyTorch’s C++ tensor library. Every operation in PyTorch ultimately maps to an ATen operator, defined with a schema string like:

mm(Tensor self, Tensor mat2) -> Tensor

The schema specifies input/output types and is used by the dispatcher, torch.fx tracing, ONNX export, and the JIT compiler. You can browse all ~2000 ATen operators at torch/_C/_VariableFunctions.pyi or via torch._C._VariableFunctions.__dir__().

Views, Copies, Contiguity, and Memory Layout¶

Views vs Copies¶

A view of a tensor shares the same underlying storage (memory buffer). Operations like reshape, view, transpose, narrow, expand, and indexing with slices typically return views:

import torch

x = torch.arange(12).reshape(3, 4)
y = x.T                  # Transpose: a view, not a copy
z = x[0:2, :]            # Slice: a view

print(x.data_ptr() == y.data_ptr())   # True -- same storage
print(x.data_ptr() == z.data_ptr())   # True -- same storage

# Modifying y modifies x
y[0, 0] = 999
print(x[0, 0])  # 999

This is critical for autograd: gradients flow through views correctly because the autograd graph records view relationships. But it also means in-place operations on views can corrupt the computation graph — PyTorch will raise a RuntimeError if you do this during a backward pass.

Memory Layout and Contiguity¶

A tensor is contiguous if its elements are laid out in row-major (C-style) order: the last dimension varies fastest. Formally, for a tensor with shape \((d_0, d_1, \ldots, d_{n-1})\) and strides \((s_0, s_1, \ldots, s_{n-1})\), contiguity means \(s_k = \prod_{j=k+1}^{n-1} d_j\) for all \(k\).

x = torch.arange(6).reshape(2, 3)
print(x.is_contiguous())   # True
print(x.stride())          # (3, 1)  -- step 3 elements between rows, 1 between cols

y = x.T                    # Transpose
print(y.is_contiguous())   # False
print(y.stride())          # (1, 3)  -- stride order reversed

# Force contiguous copy (creates new storage)
z = y.contiguous()
print(z.is_contiguous())   # True

Why does contiguity matter? Most GPU kernels (and ATen CPU kernels) assume contiguous layout. Operating on non-contiguous tensors either triggers an implicit .contiguous() copy (hurting performance) or requires a strided kernel path. In tight training loops, a silent .contiguous() can add meaningful overhead.

Strides and Non-standard Layouts¶

Strides generalize contiguity. A stride-\(s\) tensor accesses element \((i, j)\) at memory offset \(i \cdot s_0 + j \cdot s_1\). This enables:

• Transposition: swap strides without moving data.
• Broadcasting: set a stride to 0 to repeat data logically without copying.
• Slicing: adjust the base pointer and reduce the size along a dimension.

# Broadcasting via zero stride
x = torch.tensor([1.0, 2.0, 3.0])
# Expand to (4, 3) without copying:
y = x.unsqueeze(0).expand(4, 3)
print(y.stride())        # (0, 1) -- stride-0 in the batch dimension
print(y.is_contiguous()) # False
# Storage is untouched: 3 floats = 12 bytes, even though the view "has" 12 elements.
print(y.untyped_storage().nbytes())   # 12  (.storage() is the deprecated spelling)

Broadcasting Semantics¶

NumPy-style broadcasting aligns shapes from the right and stretches size-1 dimensions:

PyTorch implements this via the stride-0 trick above: a size-1 dimension that gets broadcast is assigned stride 0. No data is copied. However, when autograd differentiates through a broadcast, the backward must sum the gradient over the broadcast dimensions to match the original tensor’s shape. This is done automatically by torch.Tensor.expand’s backward and by SumBackward.

a = torch.ones(3, 1, requires_grad=True)   # shape (3, 1)
b = torch.ones(3, 4)                       # shape (3, 4), no grad

c = a + b   # broadcasts a to (3, 4)
c.sum().backward()

# grad of 'a' sums over the broadcast dim:
print(a.grad)   # tensor([[4.], [4.], [4.]])  -- sum over 4 columns

Interview Corner and Practical Patterns¶

Gradient Checkpointing (Activation Recomputation)¶

For very deep networks or long transformer sequences, storing all activations for the backward pass dominates memory. torch.utils.checkpoint.checkpoint trades compute for memory by recomputing activations during the backward pass instead of storing them:

from torch.utils.checkpoint import checkpoint

def block_forward(x, layer):
    return layer(x)

# During backward, the forward of this segment will be re-run
x_out = checkpoint(block_forward, x_in, layer, use_reentrant=False)

Mechanically, checkpoint runs the segment’s forward under no_grad (so no tape is built inside it) and installs a single Function node that, when backward reaches it, re-runs the forward with grad enabled to rebuild exactly the saved tensors it needs. The tape therefore keeps only the segment boundaries instead of every intermediate — for a transformer block that turns the ~0.6 GiB of MLP activations counted earlier into a single block-input tensor — at the cost of roughly one extra forward pass of compute per checkpointed segment (on the order of 30% added step time when every block is checkpointed). It is standard practice in LLM pretraining. See Memory-Efficient Training: Checkpointing, Offloading & LoRA Math for the full analysis, including per-operator selective checkpointing, and The Pretraining Run: A Complete Single-GPU Training Loop for the capstone’s decision about when a 100M model actually needs it.

Second-Order Gradients and create_graph¶

PyTorch can differentiate through the backward pass itself by passing create_graph=True to backward() or using torch.autograd.grad. This is used for MAML (model-agnostic meta-learning), Hessian-vector products, and implicit differentiation:

x = torch.tensor(3.0, requires_grad=True)
y = x ** 3         # y = x^3, dy/dx = 3x^2

# First-order gradient
(grad_x,) = torch.autograd.grad(y, x, create_graph=True)
print(grad_x)   # tensor(27.)  -- 3 * 3^2

# Second-order gradient (differentiates grad_x w.r.t. x)
(grad2_x,) = torch.autograd.grad(grad_x, x)
print(grad2_x)  # tensor(18.)  -- 6x = 6*3

When create_graph=True, the autograd graph for the backward computation is itself tracked, enabling higher-order differentiation. This doubles (or more) the memory cost.

torch.func (functorch): Functional Transforms¶

PyTorch’s torch.func module (formerly functorch) exposes composable functional transforms over arbitrary PyTorch functions:

from torch.func import grad, vmap, jacrev

# Compute gradient of a scalar function
f = lambda x: (x ** 2).sum()
grad_f = grad(f)
print(grad_f(torch.tensor([1.0, 2.0, 3.0])))  # tensor([2., 4., 6.])

# Batched Jacobian via vmap + jacrev
def model(params, x):
    return params @ x

J = vmap(jacrev(model, argnums=1))(params_batch, x_batch)

These transforms work at the dispatcher level, inserting themselves as dispatch keys. They compose: vmap(grad(f)) gives a batched gradient function. This is the correct modern approach to Hessian-vector products and per-sample gradients (rather than looping or using create_graph).

The Full Stack: From Python Op to GPU Kernel¶

Let us trace torch.mm(a, b) for two CUDA tensors end to end:

The backward MmBackward is registered in ATen’s derivative formulas. For \(C = AB\):

These are themselves mm calls, so they go through the same dispatcher path and are executed as cuBLAS GEMMs.

torch.compile and the Dispatcher¶

torch.compile (based on TorchDynamo + TorchInductor) traces the Python bytecode, extracts a subgraph as a torch.fx.Graph, applies fusion passes, and emits an optimized kernel. It interacts with the dispatcher by inserting a CompiledFunctionBackend dispatch key. Because the dispatcher is a clean abstraction boundary, torch.compile can replace and fuse sequences of ATen ops without touching user code. See Kernel Fusion, torch.compile, CUDA Graphs & Compilers for details.

Numerical Precision Considerations in Autograd¶

Autograd inherits the numerical properties of the operations it differentiates. Two important failure modes:

Vanishing/exploding gradients occur when VJPs amplify or attenuate the gradient signal across many layers. ReLU gates reduce flow (any negative pre-activation kills the gradient path), while sigmoid and tanh saturate (their derivatives approach zero). Residual connections, layer norm, and careful initialization (Kaiming, Xavier) combat this. See Neural Networks From Scratch: MLPs & Backprop and Numerical Computing, Floating Point & Precision.

Catastrophic cancellation in backward can occur when forward activations are large and the gradient requires subtracting nearly equal numbers. Log-softmax with cross-entropy is a canonical example: naive implementation computes \(\log(\sum \exp)\) and then subtracts, losing precision. PyTorch’s F.cross_entropy uses the log-sum-exp trick to sidestep this, and the autograd graph for the fused version is more numerically stable than separately differentiating log and softmax.

Localizing NaNs in the Backward Pass: set_detect_anomaly¶

The hardest autograd bug to diagnose is a loss or gradient that becomes NaN or Inf during the backward pass while the forward pass looked perfectly finite — e.g. an operation whose value is well-defined at a point but whose derivative is not. The backward NaN surfaces far from its cause (in optimizer.step(), or as a NaN grad norm in a logging call), so the offending forward op is hidden several layers upstream.

torch.autograd.set_detect_anomaly(True) (or the scoped context manager with torch.autograd.detect_anomaly():) turns on anomaly detection, which does two things: it checks every op’s output for NaN (note: it flags NaN, not Inf), and — crucially — it records the forward-pass Python stack trace for each Function node, so that when a backward computes a NaN it raises a RuntimeError whose traceback points at the exact forward line that created the offending op.

import torch

# A forward that is finite at x = 0 but whose *gradient* is not:
#   d/dx sqrt(x) = 0.5 / sqrt(x)  ->  inf at x = 0; the product rule below then
#   multiplies that inf by x = 0, so the backward pass hits 0 * inf = nan.
# (Plain x.sqrt() alone gives an inf grad, which anomaly mode does NOT flag --
#  it only checks for NaN -- so the product is what triggers the RuntimeError.)
x = torch.zeros(3, requires_grad=True)

# Without anomaly detection you only see a NaN grad, not WHERE it came from:
y = (x * x.sqrt()).sum()   # forward: tensor(0.) -- looks fine
y.backward()
print(x.grad)              # tensor([nan, nan, nan]) -- silent, no traceback

# With anomaly detection, backward raises AT the offending forward op:
x = torch.zeros(3, requires_grad=True)
with torch.autograd.detect_anomaly():
    y = (x * x.sqrt()).sum()  # anomaly mode stores this line's stack trace
    y.backward()
# RuntimeError: Function 'SqrtBackward0' returned nan values in its 0th output.
# The traceback includes:  "y = (x * x.sqrt()).sum()"  <- the forward line to fix

The fix is typically an epsilon or a clamp on the input to the unstable op — e.g. x.clamp_min(1e-12).sqrt(), or adding eps inside the op — after which re-running under detect_anomaly() no longer raises. The same mechanism catches the more common real-world cause in LLM training: a log, sqrt, division, or pow fed a zero or negative value from an upstream overflow. set_detect_anomaly(True) is the global switch you flip once at the top of a debug run; with torch.autograd.detect_anomaly(): scopes it to a single step.

In a large pretraining run, reach for this at STEP 3 (“isolate the step”) of the loss-spike playbook in Training Stability, Loss Spikes & Debugging Large Runs: after rolling back to the pre-spike checkpoint and replaying the exact batch, wrapping that single forward+backward in detect_anomaly() pinpoints the op, complementing the forward-hook activation probe used there to find which layer’s activations blew up first.

Key Takeaways¶

Further Reading¶

• Baydin et al., “Automatic Differentiation in Machine Learning: a Survey” (2018) — the definitive academic survey of all autodiff modes.
• Paszke et al., “Automatic differentiation in PyTorch” (NIPS 2017 Autodiff Workshop) — the original PyTorch autograd paper.
• Wengert, “A simple automatic derivative evaluation program” (1964) — the original Wengert list / tape paper.
• PyTorch documentation, “Extending PyTorch” — official reference for torch.autograd.Function and the dispatcher.
• Pytorch contributor docs, “PyTorch Dispatcher internals” (E. Yang, PyTorch blog, 2021) — deep dive into the dispatcher architecture.
• Frostig et al., “Decomposing reverse-mode automatic differentiation” (2021) — the JAX / functorch perspective on composable transforms.
• PyTorch GitHub, torch/csrc/autograd/ — the C++ engine source; reading engine.cpp and function.h is the fastest way to understand execution ordering and thread pools.

Exercises¶

1. In the leaf-tensor example, y = (x * w).sum() prints y.is_leaf == False and y.grad_fn is a SumBackward0, while x.is_leaf == True. Explain in your own words what makes a tensor a leaf versus a non-leaf, and predict what y.grad would be after y.backward(). Then describe the one-line change that would make PyTorch populate y.grad.

y = (x * w).sum()
y.retain_grad()      # ask autograd to keep this non-leaf's .grad
y.backward()
print(y.grad)        # tensor(1.)  -- now populated

2. (Quantitative — VJP by hand.) Using the network and numbers from the “Tracing autograd” worked example (\(x=[1,0]^\top\), \(W_1=\begin{bmatrix}2&1\\-1&3\end{bmatrix}\), \(w_2=[0.5,0.5]^\top\), \(L=\tfrac12 z_2^2\)), suppose you change the input to \(x=[1,1]^\top\) and leave \(W_1, w_2\) unchanged. Redo the forward and backward passes by hand and report \(z_1\), \(h\), \(z_2\), \(L\), \(\bar{w}_2\), and \(\bar{W}_1\). Which ReLU gates are open now?

3. (Conceptual.) In the MyOp/custom-Function section, StableSigmoid.forward calls ctx.save_for_backward(y) on the output y, not the input x, and the comment says this “saves memory.” Explain why saving the output is sufficient here, and give one example of an activation whose backward cannot be written from the output alone.

4. (Conceptual — broadcasting backward.) In the broadcasting example, a has shape (3, 1), b has shape (3, 4), c = a + b, and a.grad comes out as [[4.], [4.], [4.]] after c.sum().backward(). Explain where the 4 comes from, and predict a.grad if instead a had shape (1, 4) (with b still (3, 4)).

a.grad  # tensor([[3., 3., 3., 3.]])  -- shape (1, 4), summed over the 3 rows

5. (Implementation.) Implement a custom autograd.Function called ClampSTE that clamps its input to the range \([lo, hi]\) in the forward pass but uses a straight-through estimator in the backward pass — i.e. the gradient passes through unchanged (as if the clamp were the identity) regardless of whether an element was clipped. Follow the chapter’s Function conventions (static methods, ctx, .apply), remember that lo/hi are non-tensor scalars, and demonstrate on x = [-2.0, 0.5, 3.0] with lo=0.0, hi=1.0 that the forward clamps but every element’s .grad is 1.

import torch
from torch.autograd import Function


class ClampSTE(Function):
    """
    Forward:  y = clamp(x, lo, hi)      (saturates outside [lo, hi])
    Backward: dy/dx = 1 everywhere      (straight-through estimator)
    """

    @staticmethod
    def forward(ctx, x: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
        # scalars -> stash as attributes, not save_for_backward
        ctx.lo = lo
        ctx.hi = hi
        return x.clamp(lo, hi)

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        # pass gradient straight through; lo, hi are non-differentiable -> None
        return grad_output, None, None


clamp_ste = ClampSTE.apply

x = torch.tensor([-2.0, 0.5, 3.0], requires_grad=True)
y = clamp_ste(x, 0.0, 1.0)
print(y)            # tensor([0.0000, 0.5000, 1.0000])  -- clamped
y.sum().backward()
print(x.grad)       # tensor([1., 1., 1.])  -- STE: identity gradient

6. (Debugging — anomaly detection.) A colleague reports that their forward pass is finite but x.grad comes out NaN, using code of the form y = (x * x.sqrt()).sum(); y.backward() with x containing a zero. Explain, using the chapter’s account of detect_anomaly, (a) why the forward is finite but the backward is NaN, (b) why anomaly detection is what surfaces the location, and © a concrete one-line fix.

y = (x * x.clamp_min(1e-12).sqrt()).sum()

7. (Quantitative — saved-tensor accounting.) Under torch.autocast(dtype=torch.bfloat16), a nn.Linear(512, 1408, bias=False) whose master weight is fp32 is applied to an activation of 65,536 tokens. (a) Which tensors does the resulting MmBackward0 node save, and how many bytes is each? (b) How does the answer change if the surrounding block is wrapped in torch.utils.checkpoint.checkpoint(..., use_reentrant=False)?