Backpropagation › PyTorch Autograd Internals
Language AI HandbookAll chapters Backpropagation › PyTorch Autograd InternalsOne section · Sections in this chapter15 / 20 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
PyTorch's autograd engine implements the same algorithm described above, but in a highly optimized and general form. Understanding how it works demystifies the .backward() call.
The Computation Graph in PyTorch
Every PyTorch tensor with requires_grad=True participates in building a dynamic computation graph. When you perform operations on such tensors, PyTorch records the operation and its inputs, building a directed acyclic graph (DAG) of Function objects.
Each Function object knows:
- What operation was performed
- What the inputs were (stored for the backward pass)
- How to compute the VJP (vector-Jacobian product) for that operation
The graph is dynamic: it is rebuilt from scratch on each forward pass. This allows Python control flow (if statements, loops) to affect the computation graph, making PyTorch's approach more flexible than TensorFlow's original static graph approach. In a static graph framework like early TensorFlow, you described the computation once and then ran it many times. The graph had to be fixed before execution began. PyTorch's approach of building the graph as you run the Python code means you can use standard Python debugging tools, inspect intermediate values, and change the network structure from one forward pass to the next.
The grad_fn attribute of each tensor stores the function that created it. When you inspect a tensor and see grad_fn=<AddmmBackward0>, that is PyTorch telling you this tensor was created by a matrix multiply and add operation, and it knows exactly how to propagate gradients back through that operation.
The Backward Pass in PyTorch
When you call loss.backward(), PyTorch:
- Starts from the loss node with initial gradient
- Traverses the computation graph in reverse topological order
- For each
Functionnode, calls itsbackward()method with the upstream gradient - The
backward()method computes VJPs and passes them to inputs - Gradients accumulate in each leaf tensor's
.gradattribute
import torch
# Build the same network in PyTorch
torch.manual_seed(42)
W1 = torch.tensor(net.layer1.W, dtype=torch.float32, requires_grad=True)
b1 = torch.tensor(net.layer1.b, dtype=torch.float32, requires_grad=True)
W2 = torch.tensor(net.layer2.W, dtype=torch.float32, requires_grad=True)
b2 = torch.tensor(net.layer2.b, dtype=torch.float32, requires_grad=True)
def pytorch_forward(x, y_true):
"""Forward pass using PyTorch autograd."""
x_t = torch.tensor(x, dtype=torch.float32)
z1 = W1 @ x_t + b1
a1 = torch.sigmoid(z1)
z2 = W2 @ a1 + b2
log_probs = torch.log_softmax(z2, dim=0)
loss = -log_probs[y_true]
return loss
# Compute loss and gradients for first sample
loss = pytorch_forward(X[0], int(y[0]))
loss.backward()Gradient comparison: NumPy backprop vs PyTorch autograd ======================================================= W1 gradient (first row): NumPy: [1.60069347e-06 1.09165423e-06] PyTorch: [1.6010891e-06 1.0919241e-06] b1 gradient: NumPy: [ 8.01663809e-07 -1.62186445e-05 2.31950617e-07 -8.91627311e-07] PyTorch: [ 8.0186197e-07 -1.6222917e-05 2.3200158e-07 -8.9185284e-07] Max absolute difference (NumPy vs PyTorch): 8.53e-08
The gradients match to floating-point precision, confirming that our manual NumPy implementation computes the same gradients as PyTorch's optimized autograd engine.
Gradient Accumulation Patterns in PyTorch
PyTorch's gradient accumulation behavior enables several useful training patterns:
Gradient accumulation for large batches. When GPU memory is too small for the desired batch size, you can split a large batch into smaller micro-batches, call backward() on each, and only call optimizer.step() after all micro-batches. The gradients accumulate across micro-batches, producing the same gradient update as if you had run a single large batch. This technique is widely used in large language model training, where batch sizes of thousands or millions of tokens are desirable but a single token sequence might be tens of thousands of characters:
# Example: simulate batch_size=4 with micro_batch_size=2
model = nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 2))
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
X_batch = torch.tensor(X[:4], dtype=torch.float32)
y_batch = torch.tensor(y[:4], dtype=torch.long)
micro_batch_size = 2
n_micro = len(X_batch) // micro_batch_size
optimizer.zero_grad()
for i in range(n_micro):
x_micro = X_batch[i * micro_batch_size : (i + 1) * micro_batch_size]
y_micro = y_batch[i * micro_batch_size : (i + 1) * micro_batch_size]
logits = model(x_micro)
loss = (
criterion(logits, y_micro) / n_micro
) # Scale loss by micro-batch count
loss.backward() # Gradients accumulate across micro-batches
optimizer.step()Gradient norms: accumulated micro-batches vs single full-batch pass Layer Accumulated Full batch Difference ---------------------------------------------------------------- 0.weight 0.747772 0.708410 3.94e-02 0.bias 0.256722 0.244077 1.26e-02 2.weight 2.140502 2.014319 1.26e-01
The gradient norms are essentially identical, confirming that gradient accumulation faithfully simulates a larger batch size. We'll explore gradient accumulation further when we discuss distributed training in later chapters.