PyTorch Autograd Deep Dive: From Dynamic Computation Graphs to Backpropagation

Summary: Autograd is the foundation of every PyTorch training loop. This post cuts to the chase with runnable experiment code: how requires_grad triggers graph construction, what happens inside backward(), why gradients accumulate, when detach()/no_grad() are mandatory, and the engineering lessons for memory and mixed precision.

0. Experiment Goals

• See how PyTorch turns one forward pass into a dynamic computation graph (DAG)
• Understand how x.grad / w.grad are computed during backward()
• Verify with a gradient-accumulation experiment that "backward does not reset grads"
• Clarify the boundaries of detach(), torch.no_grad(), and in-place operations
• Extract the rules most relevant to training engineering (memory, AMP, gradient clipping)

1. Code Block 1: A Minimal Computation Graph

import torch

# Leaf tensors: requires_grad=True makes them tracked by the graph
x = torch.tensor([2.0, 3.0], requires_grad=True)
w = torch.tensor([4.0, 5.0], requires_grad=True)

# Forward: z = sum(x * w) elementwise
z = (x * w).sum()

print("z =", z.item())          # 2*4 + 3*5 = 23
print("z.requires_grad =", z.requires_grad)   # True

Step-by-step:

• With requires_grad=True, the tensor becomes a leaf node of the graph. Leaves are the "terminals" of backpropagation—only leaf .grad gets filled.
• x * w creates an intermediate (non-leaf) tensor; .sum() creates z. While these ops execute, PyTorch records the dependency chain x → mul → sum into a DAG.
• Each node stores two kinds of info: the forward output value (needed for backward) and the backward function (MulBackward / SumBackward).
• z.requires_grad is True because some tensor in its input chain requires grad—the grad flag propagates through the graph.

2. Code Block 2: What backward Actually Does

z.backward()

print("x.grad =", x.grad)   # tensor([4., 5.])
print("w.grad =", w.grad)   # tensor([2., 3.])

How it works:

z = Σ xᵢ·wᵢ, so ∂z/∂xᵢ = wᵢ and ∂z/∂wᵢ = xᵢ.

backward() walks the DAG in reverse topological order and executes each node's backward function:

1. SumBackward: broadcasts the upstream gradient 1.0 to each summand ([1, 1]).
2. MulBackward: for x, upstream grad × w → [1·4, 1·5] = [4, 5], stored into x.grad; for w, upstream grad × x → [2, 3], stored into w.grad.
3. After leaf grads are filled, the graph is freed (default retain_graph=False).

Key insight: gradients are products of numerical chain rule, not analytic differentiation. Each op carries its own backward function; PyTorch just multiplies upstream grads through. That is why a custom torch.autograd.Function only needs forward and backward.

3. Code Block 3: Gradient Accumulation (backward does not reset)

x = torch.tensor([2.0], requires_grad=True)
w = torch.tensor([3.0], requires_grad=True)

for i in range(3):
    z = (x * w).sum()
    z.backward()          # grads are ACCUMULATED into .grad by default

print("x.grad =", x.grad)   # tensor([9.]), i.e., 3+3+3
print("w.grad =", w.grad)   # tensor([6.]), i.e., 2+2+2

Results:

• Three backward() calls accumulate x.grad from 3.0 to 9.0—Autograd does not zero gradients by default.
• This is exactly why every training loop begins with optimizer.zero_grad(); without it, grads leak across batches and updates are corrupted.
• Conversely, this is the foundation of gradient accumulation: to simulate a larger batch, deliberately skip zeroing, run several micro-batches, then optimizer.step().

Engineering note:

# Gradient accumulation to simulate a large batch
optimizer.zero_grad()
for micro_batch in loader:
    loss = model(micro_batch)
    loss = loss / accum_steps          # normalize, otherwise effective batch grows
    loss.backward()
    if (step + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

4. Code Block 4: The Boundary of detach and no_grad

x = torch.tensor([2.0], requires_grad=True)

# Case A: use x directly
z1 = (x * 3).sum()
z1.backward()
print("A. x.grad =", x.grad)          # tensor([3.])

x.grad = None

# Case B: detach cuts the gradient flow (forward value stays, grad does not pass)
y = x.detach()
z2 = (y * 3).sum()
z2.backward()
print("B. x.grad =", x.grad)          # tensor([0.]) — gradient cut

How it works:

• detach() returns a new tensor that shares the same data but has requires_grad=False and is not in the computation graph. Gradients short-circuit at it.
• Typical uses:
  • Contrastive learning / target networks: teacher branch gradients must not flow into the student branch (MoCo, BYOL, RL target Q-networks).
  • Feature freezing: extract features with a pretrained backbone, detach(), then train only the head.
• torch.no_grad() is a context-level cut: all tensors created inside are not tracked and do not build graphs. Inference (model.eval() + prediction) must use it—saves memory and time:

with torch.no_grad():
    pred = model(x)          # no graph, no activation memory retained

⚠️ The in-place trap: in-place ops like x.add_(1) on tensors that require grad destroy the old values needed for backward and raise RuntimeError: a leaf Variable that requires grad is being used in an in-place operation. Backward needs the old forward values; overwriting them invalidates the graph.

5. Code Block 5: Writing a Custom backward

import torch

class SquareAndDouble(torch.autograd.Function):
    """Custom op: forward returns 2·x², backward grad = 4·x·g"""
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return 2 * x * x

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        return grad_output * 4 * x

x = torch.tensor([3.0], requires_grad=True)
y = SquareAndDouble.apply(x)   # 2*9 = 18
y.backward()
print("y =", y.item(), "| x.grad =", x.grad.item())   # 4*3 = 12

Results: hand-written grad 12 = ∂(2x²)/∂x = 4x = 4×3, matching chain rule. This confirms "PyTorch grad = each op's backward × upstream grad" holds for custom ops too.

Engineering note: fused ops, quantization-sensitive ops, and custom CUDA kernel gradients can all hook into the PyTorch training stack via torch.autograd.Function without touching the training framework.

6. Autograd Under Mixed Precision (AMP)

scaler = torch.cuda.amp.GradScaler()

with torch.autocast(device_type="cuda", dtype=torch.float16):
    loss = model(x)

scaler.scale(loss).backward()        # scale loss first, then backward
scaler.step(optimizer)               # unscale grads, then update
scaler.update()

How it works:

• Under AMP the forward runs in FP16 (fast, memory-saving), but gradient accumulation and parameter updates should stay FP32 to avoid precision drift.
• GradScaler scales the loss (e.g., by 2¹⁶) before backward so that small gradients that would underflow in FP16 stay representable; step() unscales automatically and skips updates when Inf/NaN appears (scaler.update() maintains the dynamic scale factor).
• For Autograd, AMP only changes the forward dtype and gradient scaling—graph construction and propagation mechanics are unchanged.

7. Engineering Takeaways (copy-paste ready for team rules)

1. Always optimizer.zero_grad() before each step: Autograd accumulates by default.
2. Use torch.no_grad() for inference/eval: no graph → ~50% activation memory saved, 30-50% faster.
3. Use detach() when crossing branches: target networks, contrastive learning, feature freezing.
4. Gradient accumulation for large-batch simulation: accumulate over micro-batches + normalize loss + step on schedule.
5. AMP requires GradScaler: prevents small-gradient underflow; do not manually unscale before scaler.step.
6. Custom ops via torch.autograd.Function: pair forward/backward, ctx.save_for_backward for intermediates.
7. Keep in-place ops away from leaf tensors: on error, suspect graph corruption.

8. Further Reading

• PyTorch Autograd Notes
• PyTorch Custom Function Tutorial
• KB: 09. Deep Learning Frameworks