Autograd Implementation: Forward and Backward Passes

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: builds graph
y = x ** 2  # Tracked

# Inference: skip graph building (faster, less memory)
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: compute gradients
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: z = (x * 2)^2
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: batch processing
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: computes gradients for all batch samples
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: gradients accumulate!
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: Gradient: tensor([...])

# 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: Stop gradients for target network in RL
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

---

Related Topics

• 05 Custom Operators - Implementing backward passes
• 03 Distributed Training - All-reduce for gradients
• 00 Readme - Memory usage during backward
• 02 Decorators - @torch.jit.script