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} \in \mathbb{R}^m\) is the upstream gradient (written as a column vector, so that \(\bar{y}^\top\) is the \(1 \times m\) row vector that multiplies \(J\) from the left) and \(\bar{x} \in \mathbb{R}^n\) 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} = J^\top \bar{y}\). 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:

accumulation_steps = 8         # number of micro-batches per optimizer step

batches = iter(loader)         # a DataLoader is iterable, NOT an iterator:
                               # next(loader) raises TypeError -- call iter() first

optimizer.zero_grad()          # clear accumulated grads
for _ in range(accumulation_steps):
    inputs, targets = next(batches)
    loss = loss_fn(model(inputs), targets) / 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 is the module-granularity analogue of tensor.register_hook, not of this one: it sees the gradients w.r.t. the module’s inputs and outputs, and it fires before that module’s parameter AccumulateGrad nodes have run — inside such a hook module.weight.grad is still None (or stale). Use register_post_accumulate_grad_hook when you need “this parameter’s gradient is ready.”

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 what keeps the all-reduce from serializing behind the backward pass. It hides communication time; it does not reduce the volume moved, so a run whose gradient bytes divided by link bandwidth exceed the backward’s compute time is still bandwidth-bound — which is what motivates bf16 all-reduce, gradient compression, and ZeRO/FSDP sharding. 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 expression \(\sigma(x) = 1/(1+e^{-x})\) has a well-behaved forward: for large negative \(x\), \(e^{-x}\) overflows to inf and 1/(1+inf) rounds to exactly 0.0; for large positive \(x\), \(e^{-x}\) underflows to 0 and the result is exactly 1.0. Both are correctly rounded. The damage shows up elsewhere. First, if you let autograd differentiate that expression, the overflowed intermediate poisons the backward: torch.tensor([-1000.], requires_grad=True) run through 1/(1+torch.exp(-x)) yields a nan gradient (the chain rule multiplies a \(0\) by an \(\infty\)), whereas torch.sigmoid returns \(0\). Second, saturating to exactly \(0\) or \(1\) destroys every downstream bit of \(\log \sigma(x)\) — the reason PyTorch ships logsigmoid and binary_cross_entropy_with_logits. So in practice you call torch.sigmoid, which is what we do inside the forward below; the point of the example is the custom-backward interface, and a hand-written analytic backward that reads the output is also what keeps the gradient finite in the saturated regime:

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.8236 0.4272 0.1017 0.6384]
# grad   : [0.1453 0.2447 0.0913 0.2308]  -- 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: vmap and grad (torch.func) work by inserting themselves at a dispatch key layer, intercepting operations; torch.compile’s tracing layer rides the same mechanism via __torch_dispatch__ modes (its frontend, TorchDynamo, is a bytecode hook rather than a key — see below).
• 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. This is also the precise difference between the two inference contexts. no_grad flips a thread-local GradMode flag: the Autograd kernel still runs, but it skips allocating a node and immediately redispatches. inference_mode goes further and actually excludes the autograd keys from the thread-local dispatch set (you can see this with torch._C._dispatch_tls_local_exclude_set()), so the call lands straight on the backend kernel — which is where its extra speed over no_grad comes from.

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. An implicit broadcast inside an elementwise op does not insert any node to do this — inspect c.grad_fn below and you will find a bare AddBackward0 whose next_functions points straight at a’s AccumulateGrad. The reduction is performed by the engine itself: validate_outputs in torch/csrc/autograd/engine.cpp compares each produced gradient against the recorded shape of the input edge and calls at::sum_to when they differ but are expandable. (An explicit x.expand(...) does get an ExpandBackward0 node, whose backward performs the same summation.)

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, the modern use_reentrant=False path runs the segment’s forward with grad enabled, but wraps it in a saved_tensors_hooks pair whose pack returns an empty placeholder instead of the tensor. So the graph nodes inside the segment are built exactly as usual (out.grad_fn really is a ReluBackward0, and its next_functions really do point at the matmul below it) — they simply hold no activations. When backward reaches one of them, unpack re-runs the segment’s forward to regenerate the tensor it asked for. (The older, deprecated use_reentrant=True path is the one that runs the segment under no_grad and installs a single opaque Function node.) Either way the tape keeps only the segment boundaries worth of data instead of every intermediate — for a transformer block that turns the ~0.75 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

params_batch = torch.randn(8, 3, 4)   # 8 examples, each a (3, 4) weight
x_batch      = torch.randn(8, 4)      # 8 inputs of size 4

J = vmap(jacrev(model, argnums=1))(params_batch, x_batch)
print(J.shape)   # torch.Size([8, 3, 4]) -- one (3, 4) Jacobian d(out)/dx per example

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. Note that TorchDynamo itself is not a dispatch key: it hooks CPython’s frame evaluation to capture bytecode. The dispatcher-level machinery sits downstream — AOTAutograd lowers the captured graph to ATen using __torch_dispatch__ modes (riding the Python and Functionalize keys), and the compiled forward/backward pair is installed into the tape as an ordinary autograd.Function node. 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)?