Part 3

Neural networks
in PyTorch

The previous part explained the idea. This one is the code — the library every WillMe model is built with, in the four pieces you actually need: a tensor, a module, an autograd graph, and a loop.

Everything is a tensor

A tensor is an array of numbers with a shape. A single number has shape (), a list has shape (5,), a batch of token embeddings has shape (batch, tokens, 768). Most bugs in model code are shape bugs, which is why the shape is the first thing anyone prints.

Two things make a tensor more than an array: it can live on a GPU, and it can remember how it was computed. That second one is what makes training possible at all.

A tensor moving through one layer

Shape

(1, 12)

Pick a stage to see what the shape becomes and why. These are WillMe GPT 3.3's real dimensions: 768 channels, 12 heads of 64, a 2,048-token window.

A model is a class with a forward method nn.Module tracks every parameter you assign to it, so .parameters(), .to(device) and .state_dict() all work without you keeping a list.

PyTorch models subclass nn.Module. You declare the layers in __init__ and describe how data flows through them in forward. There is no separate backward method to write: PyTorch derives it from what forward did.

Below is a two-layer network — the same shape as the diagram on the previous part — alongside a real block from this project.

The same network, in code
import torch
import torch.nn as nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(3, 5)   # 3×5 weights + 5 biases
        self.out    = nn.Linear(5, 2)   # 5×2 weights + 2 biases

    def forward(self, x):
        x = torch.relu(self.hidden(x))  # the squashing step
        return self.out(x)

model = TinyNet()
print(sum(p.numel() for p in model.parameters()))  # 27

The left tab is the toy from part 2. The right tabs are real files in this repository: an RWKV block from v3_model.py and the GRU decoder that serves WillMe GRU 2.

nn.Linear is the neuron from part 2, in bulk

nn.Linear(3, 5) is five neurons that each read the same three inputs. Internally it holds one weight matrix of shape (5, 3) and one bias vector of shape (5,), and computes x @ W.T + b for the whole batch at once. That is 15 + 5 = 20 parameters.

Every projection inside WillMe GPT 3.3 is one of these with the bias switched off. The query projection is nn.Linear(768, 768, bias=False) — 589,824 parameters in a single call.

Autograd: the tape that remembers Every operation on a tensor that requires grad records itself. loss.backward() walks that record backwards, filling in .grad on every parameter.

In part 2 the gradient of the line fit was worked out by hand, which was fine for two parameters. For a hundred million, PyTorch does it. As the forward pass runs, each operation quietly records what it did. Calling loss.backward() walks that recording in reverse, applying the chain rule, and leaves a gradient sitting on every parameter.

Watch the tape build and unwind
Ready

Step through a forward and backward pass on a two-parameter model. The gradients shown are the true derivatives for this expression, computed the same way autograd would.

The loop that does the work

Every training run in this project, from the 15M-parameter GRU 1 to GPT 3.3, is the same five lines repeated. Hover the numbers for what each line is for.

The five lines
1  optimizer.zero_grad()2  logits = model(batch)3  loss = loss_fn(logits, targets)4  loss.backward()5  optimizer.step()

Miss line 1 and gradients from the last batch are added to this one — the single most common PyTorch bug, and a silent one.

The same loop, as WillMe GPT 3.3 actually runs it

Four differences, all of them about fitting a real model onto real hardware rather than about the idea: 

with torch.autocast(device_type=device, dtype=torch.bfloat16):
    logits = model(batch)
    loss = loss_fn(logits.float(), targets)   # loss in FP32

(loss / accum_steps).backward()               # accumulate 16 micro-batches

if step % accum_steps == 0:
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    scheduler.step()                          # cosine decay from 4e-4
    optimizer.zero_grad(set_to_none=True)

Autocast runs the forward pass in bfloat16 for speed while the loss is computed in FP32 for stability. Gradient accumulation makes 16 small batches behave like one large one, because the whole batch will not fit in memory. Clipping caps the gradient norm at 1.0 so a single strange batch cannot blow the weights apart. And the scheduler decays the learning rate along a cosine curve over all 61,035 updates.

Where the numbers actually live

A trained model is just its parameters, and PyTorch stores them in a state_dict — a plain dictionary from parameter name to tensor. Saving is torch.save(model.state_dict(), path) and loading it back rebuilds the architecture in code first, then pours the numbers in. The .pth files in this project's models/ folder are exactly that, which is why they are useless without the matching class.

Two habits that save hours

Call model.eval() before inference. Dropout and batch norm behave differently while training, and forgetting the switch makes a working model produce quietly worse output. This project's engine.py does it right after loading each checkpoint.

Wrap inference in torch.no_grad(). Without it PyTorch builds the autograd tape for text you are never going to backpropagate through, and the memory goes with it.

That is the whole toolkit

Tensors carry the numbers, modules hold the parameters, autograd finds the gradients, and a five-line loop applies them. Everything from here is a question of what shape you build out of those pieces and what data you point it at. The next part starts on the first of those: what turns a network into a language model.