9. From-Scratch MLP Training Step: Forward, Loss, Backward, Optimizer

Problem

Before any RLHF-specific machinery, an interviewer can just check whether you can write the four-line skeleton every PyTorch training loop is built from, in the right order, with no shortcuts: zero out stale gradients, run the forward pass, compute the loss, backpropagate, and step the optimizer.

Getting the order wrong — stepping before backward, or forgetting to zero gradients — silently corrupts training without raising an error, which is exactly what makes it worth testing directly rather than assuming it's obvious.

Source: src/9_mlp_training_step.py

def train_step(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    x: torch.Tensor,
    y: torch.Tensor,
    loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
) -> torch.Tensor: ...

>>> model = nn.Linear(2, 1)
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
>>> x = torch.tensor([[1.0, -1.0]])
>>> y = torch.tensor([[1.0]])
>>> loss = train_step(model, optimizer, x, y, nn.MSELoss())
>>> loss.requires_grad
False

Step-by-Step Approach

1. Call optimizer.zero_grad() first — gradients accumulate into .grad by default in PyTorch, so skipping this adds the new step's gradient on top of whatever was left over from the previous step.
2. Run the forward pass: predictions = model(x).
3. Compute the loss: loss = loss_fn(predictions, y).
4. Call loss.backward(), which populates .grad on every parameter with requires_grad=True that contributed to the loss.
5. Call optimizer.step() — this must come strictly after backward(), since it reads the just-populated .grad values to update the parameters.
6. Return loss.detach(), not the raw loss — the caller almost never needs the still-attached computation graph, and holding onto it needlessly keeps memory allocated.

The key insight is that this exact five-line sequence — zero, forward, loss, backward, step — is the one piece of PyTorch boilerplate that has to be memorized cold, because every bug in the wrong order (stepping before backward, forgetting to zero) trains a model that runs without any error and just quietly doesn't learn correctly.

Reference solution

import torch
import torch.nn as nn


def train_step(model, optimizer, x, y, loss_fn):
    # zero_grad first: gradients accumulate by default in PyTorch, so
    # skipping this adds the new step's gradient on top of the last one
    optimizer.zero_grad()
    predictions = model(x)
    loss = loss_fn(predictions, y)
    loss.backward()  # populates .grad on every parameter with requires_grad=True
    optimizer.step()  # must come AFTER backward -- it reads .grad to update params
    return loss.detach()  # detach: caller gets the number, not a live graph handle

Key Functions & Tricks

• optimizer.zero_grad() — clears accumulated .grad on every parameter the optimizer manages; call it before backward(), not after.
• loss.backward() — walks the autograd graph backward from loss, accumulating gradients into every leaf tensor with requires_grad=True along the way.
• optimizer.step() — applies one update to every managed parameter using its current .grad; a no-op if .grad is still None or stale.
• loss.detach() — returns the same numeric value without the graph attached, the correct thing to hand back to a caller that just wants to log or compare the loss.

How to Recognize This Pattern