Skip to content

Autograd Implementation

Overview

Automatic differentiation (autograd) is PyTorch's core innovation:

  • Forward pass: Compute outputs and build computation graph
  • Backward pass: Compute gradients by traversing graph
  • Custom functions: Extend autograd for custom operations
  • Gradient tracking: requires_grad, grad_fn

Understanding autograd reveals how to build ML frameworks.

Core Concepts: Computation Graphs

Tensor Computation Graph

import torch

# Tensors track operations
x = torch.tensor([2.0], requires_grad=True)
y = torch.tensor([3.0], requires_grad=True)

# Operations create nodes in computation graph
z = x * y # Multiply operation node
w = z + x # Add operation node
loss = w.sum() # Sum operation node

# Graph structure (computational dependency):
# x, y (leaf nodes)
 - # │ │
 - # └──* (MulBackward)
 - # │
- # +──+ (AddBackward)
 - # │
# sum (SumBackward)
 - # │
# loss

print(loss.grad_fn) # <SumBackward0>
print(z.grad_fn) # <MulBackward0>
print(x.grad_fn) # None (leaf tensor)

grad_fn Chain

Each intermediate tensor remembers its operation:

import torch

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

z = x ** 2
print(z.grad_fn) # <PowBackward0>

w = z + 5
print(w.grad_fn) # <AddBackward0>

# Each node remembers its inputs and operation
print(z.grad_fn.next_functions) # References to previous grad_fn

# Follow the chain
print(w.grad_fn.next_functions[0]) # Points to z's grad_fn
print(w.grad_fn.next_functions[1]) # None (5 is constant)

Forward Pass: Building the Graph

Tracing Operations

import torch

def forward(x):
 """Simple forward pass."""
 y = x ** 2 # PowBackward
 z = y + 3 # AddBackward
 w = z * 2 # MulBackward
 return w.sum() # SumBackward

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
loss = forward(x)

# At this point, computation graph is built
# But no gradients computed yet
print(x.grad) # None

Tracking with requires_grad

import torch

# Only tensors with requires_grad=True are tracked
x = torch.tensor([1.0], requires_grad=True)
y = torch.tensor([2.0], requires_grad=False) # Won't track

z = x * y # z inherits requires_grad from x
print(z.requires_grad) # True
print(z.grad_fn) # <MulBackward>

# Tensors without requires_grad don't build graph
z2 = x * 10 # x has requires_grad=True
print(z2.grad_fn) # <MulBackward>

# Detach stops tracking
z3 = x.detach() * y
print(z3.requires_grad) # False
print(z3.grad_fn) # None

No-Grad Context: Skip Graph Building

import torch

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

# Normal
y = x ** 2 # Tracked

# Inference
with torch.no_grad():
 z = x ** 2 # Not tracked
 print(z.requires_grad) # False
 print(z.grad_fn) # None

# Typical use
model.eval()
with torch.no_grad():
 predictions = model(test_data) # Much faster

-

Backward Pass: Computing Gradients

Backward Propagation

import torch

# Forward
x = torch.tensor([2.0, 3.0](/2.0,-3.0/), requires_grad=True)
W = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)

y = x @ W # Matrix multiply
loss = y.sum() # Sum all elements

# Backward
loss.backward()

# Gradients computed
print(x.grad) # Gradient w.r.t. input
print(W.grad) # Gradient w.r.t. weights
print(loss.grad) # None (loss is scalar)

Manual Gradient Computation

To understand what backward() does:

import torch

# Forward
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
loss = y

# Manual backward computation
# loss = y = x^2
# ∂loss/∂x = 2x = 2 * 2 = 4

loss.backward()
print(x.grad) # tensor([4.])

# Verify
print(2 * x.data) # tensor([4.])

Chain Rule Implementation

PyTorch uses chain rule in backward pass:

import torch

# forward
x = torch.tensor([3.0], requires_grad=True)

y = x * 2 # First operation
z = y ** 2 # Second operation

# Backward uses chain rule:
# dz/dx = dz/dy * dy/dx
# dz/dx = 2*y * 2 = 4*y = 4*6 = 24

z.backward()
print(x.grad) # tensor([24.])

# Manual verification
# y = 3 * 2 = 6
# z = 6^2 = 36
# dz/dy = 2 * 6 = 12
# dy/dx = 2
# dz/dx = 12 * 2 = 24

Batch Backward

import torch

# Typical
x = torch.randn(4, 10, requires_grad=True) # Batch of 4
target = torch.randn(4, 5)

model = torch.nn.Linear(10, 5)

# Forward
output = model(x)
loss = torch.nn.functional.mse_loss(output, target)

# Backward
loss.backward()

# Gradients accumulated for entire batch
print(model.weight.grad.shape) # (5, 10) - same as weight shape
print(x.grad.shape) # (4, 10) - same as input shape

-

Custom Autograd Functions

Implementing Custom Forward/Backward

import torch

class ReLU(torch.autograd.Function):
 """Custom ReLU implementation."""

 @staticmethod
 def forward(ctx, input):
 """
 Compute ReLU: max(0, x)

 ctx: Context object to save information for backward
 """
 # Save input for backward
 ctx.save_for_backward(input)

 # Forward computation
 return torch.clamp(input, min=0)

 @staticmethod
 def backward(ctx, grad_output):
 """
 Compute gradient of ReLU

 grad_output: gradient from next layer (dL/dy)
 returns: gradient w.r.t. input (dL/dx)
 """
 # Retrieve saved tensors
 input, = ctx.saved_tensors

 # ReLU gradient: 1 if x > 0, else 0
 grad_input = grad_output.clone()
 grad_input[input < 0] = 0

 return grad_input

# Use custom function
def relu(x):
 return ReLU.apply(x)

# Test
x = torch.tensor([1.0, -2.0, 3.0], requires_grad=True)
y = relu(x)
loss = y.sum()
loss.backward()

print(x.grad) # [1, 0, 1] - gradient only where x > 0

Complex Custom Function: Matrix Operation

import torch

class LinearWithCustomGrad(torch.autograd.Function):
 """Custom linear layer with modified gradient."""

 @staticmethod
 def forward(ctx, input, weight, bias=None):
 # Save for backward
 ctx.save_for_backward(input, weight, bias)

 # Forward: y = x @ W^T + b
 output = torch.matmul(input, weight.t())
 if bias is not None:
 output += bias

 return output

 @staticmethod
 def backward(ctx, grad_output):
 # Retrieve saved tensors
 input, weight, bias = ctx.saved_tensors

 # Compute gradients
 # ∂L/∂input = ∂L/∂output @ weight
 grad_input = torch.matmul(grad_output, weight)

 # ∂L/∂weight = ∂L/∂output^T @ input
 grad_weight = torch.matmul(grad_output.t(), input)

 # ∂L/∂bias = sum of ∂L/∂output
 grad_bias = grad_output.sum(0) if bias is not None else None

 # Return gradients in same order as forward arguments
 return grad_input, grad_weight, grad_bias

# Usage
linear = LinearWithCustomGrad()

x = torch.randn(32, 10, requires_grad=True)
W = torch.randn(5, 10, requires_grad=True)
b = torch.randn(5, requires_grad=True)

y = linear.apply(x, W, b)
loss = y.sum()
loss.backward()

print(x.grad.shape) # (32, 10)
print(W.grad.shape) # (5, 10)
print(b.grad.shape) # (5,)

-

Gradient Accumulation

Default: Accumulation Behavior

import torch

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

# First backward
y = x ** 2
y.backward()
print(x.grad) # tensor([2.])

# Second backward
y = x ** 3
y.backward()
print(x.grad) # tensor([2. + 3.]) = tensor([5.])

# Why? Useful for:
# - Multi-task learning
# - Gradient accumulation across batches
# - Meta-learning

Zero Gradients Between Updates

import torch

optimizer = torch.optim.SGD([x], lr=0.01)
model = torch.nn.Linear(10, 5)

for epoch in range(100):
 # Forward
 output = model(input_data)
 loss = criterion(output, target)

 # Backward accumulates gradients
 loss.backward()

 # Must zero gradients before next iteration!
 optimizer.zero_grad() # or model.zero_grad()

 # Update parameters
 optimizer.step()

# If you forget zero_grad, gradients keep accumulating!

Intentional Gradient Accumulation

import torch

# Simulate large batch by accumulating gradients
model = torch.nn.Linear(10, 5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

effective_batch_size = 128
actual_batch_size = 32
accumulation_steps = effective_batch_size // actual_batch_size

for epoch in range(100):
 for batch_idx, (x, y) in enumerate(train_loader):
 # Forward
 pred = model(x)
 loss = criterion(pred, y) / accumulation_steps

 # Backward: gradients accumulate
 loss.backward()

 # Update after K batches
 if (batch_idx + 1) % accumulation_steps == 0:
 optimizer.step()
 optimizer.zero_grad()

-

Gradient Hooks: Monitoring/Modifying Gradients

Register Hooks

import torch

def print_grad_hook(grad):
 """Hook that prints gradient."""
 print(f"Gradient: {grad}")
 return grad

x = torch.randn(3, requires_grad=True)
hook_handle = x.register_hook(print_grad_hook)

y = (x ** 2).sum()
y.backward()
# Output

# Unregister hook if needed
hook_handle.remove()

Gradient Clipping Example

import torch

def clip_grad_hook(grad, max_norm=1.0):
 """Hook that clips gradient norm."""
 norm = grad.norm()
 if norm > max_norm:
 return grad * (max_norm / norm)
 return grad

# Register for all parameters
model = torch.nn.Linear(10, 5)
for name, param in model.named_parameters():
 param.register_hook(lambda g, m=max_norm: clip_grad_hook(g, m))

# Training
y = model(x)
loss = y.sum()
loss.backward() # Gradients automatically clipped

Monitor Parameter Gradients

import torch

class GradientMonitor:
 def __init__(self, model):
 self.model = model
 self.grad_stats = {}
 self._register_hooks()

 def _register_hooks(self):
 """Register hooks on all parameters."""
 for name, param in self.model.named_parameters():
 param.register_hook(lambda g, n=name: self._record_grad(n, g))

 def _record_grad(self, name, grad):
 """Record gradient statistics."""
 self.grad_stats[name] = {
 'mean': grad.mean().item(),
 'std': grad.std().item(),
 'min': grad.min().item(),
 'max': grad.max().item(),
 }
 return grad

 def print_stats(self):
 """Print gradient statistics."""
 for name, stats in self.grad_stats.items():
 print(f"{name}: mean={stats['mean']:.6f}, std={stats['std']:.6f}")

# Usage
model = torch.nn.Sequential(
 torch.nn.Linear(10, 20),
 torch.nn.ReLU(),
 torch.nn.Linear(20, 5)
)

monitor = GradientMonitor(model)

# Training step
y = model(x)
loss = y.sum()
loss.backward()

monitor.print_stats()

Detaching and No-Grad

When to Detach

import torch

x = torch.randn(3, 4, requires_grad=True)

# Detach breaks gradient flow
y = x ** 2
z = y.detach() # Stop tracking

w = z ** 2
loss = w.sum()
loss.backward()

print(x.grad) # None - gradient didn't flow through detach
print(y.grad) # None - y doesn't have grad_fn anymore

# Use case
target_output = target_network(state).detach()

No-Grad for Inference

import torch

model = torch.nn.Linear(10, 5)

# Training
x = torch.randn(3, 10)
y = model(x)
loss = y.sum()
loss.backward() # Computes gradients, builds graph

# Inference
with torch.no_grad():
 predictions = model(x)
 # No graph building, much faster, less memory

-

Gradient Checking

Numerical Gradient Verification

import torch

def numerical_gradient(func, x, eps=1e-5):
 """Compute numerical gradient using finite differences."""
 grad = torch.zeros_like(x)

 for i in range(x.numel()):
 x_plus = x.clone()
 x_plus.view(-1)[i] += eps

 x_minus = x.clone()
 x_minus.view(-1)[i] -= eps

 grad.view(-1)[i] = (func(x_plus) - func(x_minus)) / (2 * eps)

 return grad

# Test custom backward implementation
def my_function(x):
 return (x ** 2).sum()

x = torch.randn(3, requires_grad=True)

# Analytical gradient (via autograd)
y = my_function(x)
y.backward()
analytical_grad = x.grad.clone()

# Numerical gradient
x.grad.zero_()
numerical_grad = numerical_gradient(my_function, x)

# Compare
print(torch.allclose(analytical_grad, numerical_grad, atol=1e-4))

Practical: Training Loop with Autograd

import torch
import torch.nn.functional as F

class SimpleNet(torch.nn.Module):
 def __init__(self):
 super().__init__()
 self.fc1 = torch.nn.Linear(10, 20)
 self.fc2 = torch.nn.Linear(20, 5)

 def forward(self, x):
 x = F.relu(self.fc1(x))
 return self.fc2(x)

# Setup
model = SimpleNet()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()

# Training loop
for epoch in range(10):
 for batch_x, batch_y in train_loader:
 # Forward pass: build computation graph
 predictions = model(batch_x)

 # Compute loss: creates loss node in graph
 loss = criterion(predictions, batch_y)

 # Backward pass: traverse graph, compute gradients
 optimizer.zero_grad() # Clear old gradients
 loss.backward() # Compute new gradients

 # Update: modify parameters using gradients
 optimizer.step()

# Evaluation (no gradient computation)
model.eval()
with torch.no_grad():
 for batch_x, batch_y in test_loader:
 predictions = model(batch_x)
 accuracy = (predictions.argmax(1) == batch_y).float().mean()

Summary: Autograd Flow

1. FORWARD PASS
 x → [Op1] → y → [Op2] → z → [Op3] → loss

 Builds computation graph, saving:
 - Operation type
 - Input tensors
 - grad_fn pointers

2. BACKWARD PASS (loss.backward())
 loss (grad = 1)
 ↓
 [Op3.backward] → grad_z
 ↓
 [Op2.backward] → grad_y
 ↓
 [Op1.backward] → grad_x
 ↑
 Leaf tensors (.grad attribute)

3. GRADIENT UPDATE
 param = param - lr * param.grad

-