Autograd Internals

Autograd is the engine that makes PyTorch a deep learning framework rather than a fast tensor library. Every operation performed on a tensor that requires gradients is recorded in a dynamic computational graph; calling loss.backward() then walks that graph in reverse, applying the chain rule to compute gradients of the loss with respect to every leaf parameter. Without autograd, every gradient in every model would have to be derived and coded by hand; with autograd, the framework does that work automatically, every time.

This section explains how the engine works under the hood, when to use the various ways of controlling it, and how to extend it with custom operations and hooks. The mathematical justification for the chain rule itself is covered in Appendix A.3 (Calculus for Machine Learning); here the focus is on the mechanics PyTorch exposes.

The Computational Graph

A tensor that participates in a differentiable computation carries two pieces of extra state. The first is the boolean flag requires_grad; if true, autograd will track operations on this tensor. The second is the grad_fn attribute, which points to a node in the dynamic graph that records the operation that produced this tensor and holds references to the upstream tensors needed to compute the local gradient.

Leaf tensors (model parameters, typically) have requires_grad=True and grad_fn=None because no operation produced them. Intermediate tensors have requires_grad=True only if at least one of their inputs did, and their grad_fn records the producing op (for example AddmmBackward0 for the result of a linear layer). The graph is built on the fly during the forward pass and discarded after each backward pass; that is what "dynamic computation graph" means.

import torch

w = torch.tensor([1.5], requires_grad=True)       # leaf parameter
b = torch.tensor([0.3], requires_grad=True)       # leaf parameter
x = torch.tensor([2.0])                            # leaf data, no grad

z = x * w + b                                       # intermediate
a = torch.sigmoid(z)                                # intermediate
loss = (a - 1.0) ** 2                               # scalar output

print("w.grad_fn:", w.grad_fn)
print("z.grad_fn:", z.grad_fn)
print("a.grad_fn:", a.grad_fn)
print("loss.grad_fn:", loss.grad_fn)

backward() and Gradient Accumulation

Calling loss.backward() traverses the graph from loss back to every leaf with requires_grad=True, computing partial derivatives by the chain rule and writing them into the leaf's .grad attribute. A scalar loss requires no arguments. A non-scalar output requires a gradient= argument of the same shape, which is multiplied element-wise with the implicit upstream gradient (typically used for advanced cases like Jacobian-vector products).

A subtle and important behavior: .grad accumulates. Each call to backward() adds the new gradient to whatever is already in .grad. This is by design; it enables gradient accumulation across multiple forward passes (useful for simulating a larger batch). It also means that in a normal training loop, optimizer.zero_grad() must be called before every backward pass to clear stale gradients. Forgetting this is the most common training-loop bug.

The other way to compute gradients is torch.autograd.grad(outputs, inputs), which returns the gradients directly without writing them to .grad. This is useful when computing gradients with respect to specific tensors (rather than all parameters) or when implementing meta-learning, second-order methods, or implicit-differentiation tricks. By default grad() consumes the graph; pass retain_graph=True to keep it for additional gradient calls, and create_graph=True to enable differentiation of the gradient itself (for Hessians and other higher-order quantities).

Controlling Gradient Tracking

Gradient tracking costs both memory (intermediate activations must be retained for the backward pass) and compute (every op records its backward function). Disabling tracking when it is not needed is a major optimization, especially during inference and validation. PyTorch offers three mechanisms for this, each with a different scope.

The first is tensor.detach(): a method that returns a new tensor sharing storage with the original but with requires_grad=False and no grad_fn. The original is unchanged. detach() is the surgical tool: it removes one specific tensor from the graph. The canonical use is for logging, where loss.detach() records the scalar value without retaining the backward graph.

The second is the context manager torch.no_grad(): every tensor produced inside the with block has requires_grad=False regardless of its inputs. This is the standard wrapper for inference and validation passes; it cuts memory use roughly in half because no activations are kept around.

import torch

model = torch.nn.Linear(50, 3)
X = torch.rand(8, 50)

# Inference: do not build the graph.
model.eval()
with torch.no_grad():
    logits = model(X)
    probs = torch.softmax(logits, dim=1)

print(logits.requires_grad, logits.grad_fn)

The third is torch.set_grad_enabled(bool): a context manager (or function) that toggles tracking globally inside its scope. The form with torch.set_grad_enabled(is_train): is useful inside a function that handles both training and evaluation, so the same loop body can be reused for both.

A related tool is torch.inference_mode(), a newer and stricter version of no_grad() that disables view tracking and a few other autograd bookkeeping mechanisms for extra speed. Use it for pure inference workloads where no tensor produced inside the block will ever participate in training; otherwise stick with no_grad().

Custom Autograd Functions

The full power of autograd is exposed through torch.autograd.Function, the base class that defines what it means to be a differentiable operation. Subclassing it requires implementing two static methods: forward(ctx, *inputs) computes the forward output and stashes anything needed for the backward pass on ctx; backward(ctx, *grad_outputs) retrieves the saved tensors and computes the gradients of each input.

import torch

class ScaledReLU(torch.autograd.Function):
    """ReLU multiplied by a scalar `alpha`, with custom backward."""

    @staticmethod
    def forward(ctx, x, alpha):
        ctx.save_for_backward(x)
        ctx.alpha = alpha
        return alpha * torch.relu(x)

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        grad_input = grad_output * ctx.alpha * (x > 0).to(grad_output.dtype)
        # Inputs: x, alpha. alpha is not a Tensor, so return None for it.
        return grad_input, None


x = torch.randn(4, requires_grad=True)
y = ScaledReLU.apply(x, 2.0).sum()
y.backward()
print(x.grad)

from torch.func import grad, vmap

def loss_fn(params, x, y):
    return ((params @ x - y) ** 2).sum()

# Per-sample gradients: vmap the single-sample grad over the batch axis.
per_sample_grad = vmap(grad(loss_fn), in_dims=(None, 0, 0))(params, X, y)

Hooks: Inspecting and Modifying the Graph

Hooks are callbacks that fire during the forward or backward pass without changing the model's source code. They are the diagnostic tool of choice for inspecting activations, gradients, and parameter updates in a running model. Three types exist on nn.Module: register_forward_pre_hook (fires before forward), register_forward_hook (fires after forward, sees inputs and outputs), and register_full_backward_hook (fires during backward, sees gradient inputs and outputs). Each returns a handle whose .remove() method detaches the hook.

Tensor-level hooks are even finer-grained. tensor.register_hook(fn) calls fn(grad) every time a gradient flows through that tensor and lets fn return a modified gradient (or None to leave it untouched). This is the right level for gradient clipping or noise injection on a specific intermediate.

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 4))

activations = {}
def save_act(name):
    def hook(module, inputs, output):
        activations[name] = output.detach()
    return hook

handles = [
    model[0].register_forward_hook(save_act("linear_0")),
    model[2].register_forward_hook(save_act("linear_1")),
]

_ = model(torch.randn(2, 8))
for k, v in activations.items():
    print(k, tuple(v.shape))

for h in handles:
    h.remove()

Anomaly Detection

When a backward pass produces NaN or Inf gradients, the most useful debugging aid is torch.autograd.set_detect_anomaly(True). With anomaly detection on, PyTorch records the Python stack trace of every forward op and replays it whenever a non-finite gradient is detected during the backward pass, so the offending operation can be identified precisely. The cost is significant slowdown, so this is a debugging tool rather than a production setting. Section E.9 walks through the full NaN-hunting workflow.