Lesson 5 - Guided Project: Train Your Mini-GPT

Welcome to the Guided Project#

This is the payoff. Over six modules you built every piece of a transformer and proved each one correct; over this module you built the machinery to train it. Now you’ll put it all in motion: configure a model, run a full training session, watch the loss fall while checking it isn’t overfitting, see the model produce recognizable text, and save the trained weights to disk. What you finish with is a real, trained language model — small, but genuinely working — ready to generate from in Module 8.

By the end of this project, you will be able to:

• Configure and instantiate a mini-GPT and report its parameter count
• Run a complete training loop while monitoring training and validation loss
• Do a greedy continuation and judge what the model has learned from its output
• Save the trained parameters to disk and reload them to reproduce the model exactly

Every stage below runs for real in pure NumPy, reusing the MiniGPT, get_batch, and Adam you built. Let’s train it.

Stage 1: Configure the Model#

We choose a deliberately small configuration so the whole run finishes in seconds on a CPU, then tokenize the corpus and hold out a validation split (Lesson 1).

import numpy as np, time
# MiniGPT (Module 6), get_batch + tokenized data/vocab_size/decode (Lesson 1),
# Adam (Lesson 2) are reused unchanged.

block_size, n_embd, n_head, n_layer = 32, 64, 4, 3
batch_size, lr, steps = 32, 3e-3, 1200

n = int(len(data) * 0.9)
train_data, val_data = data[:n], data[n:]

model = MiniGPT(vocab_size, block_size, n_embd, n_head, n_layer, seed=42)
opt = Adam(model.P, lr=lr)
print("config: n_embd=%d n_head=%d n_layer=%d block_size=%d" % (n_embd, n_head, n_layer, block_size))
print("parameters:", model.n_params())
print("train tokens:", len(train_data), "| val tokens:", len(val_data))

config: n_embd=64 n_head=4 n_layer=3 block_size=32
parameters: 154456
train tokens: 9018 | val tokens: 1002

A three-layer, four-head model with a 64-dimensional embedding comes to 154,456 parameters — tiny by modern standards, plenty to learn a ten-sentence corpus.

Stage 2: Train, and Watch Both Losses#

Now the training loop from Lesson 3, with a validation estimate (Lesson 4) every 200 steps so we can confirm the model is learning rather than memorizing.

def estimate_loss(model, data, rng, iters=20):
    total = 0.0
    for _ in range(iters):
        xb, yb = get_batch(data, block_size, batch_size, rng)
        _, loss = model.forward(xb, yb)
        total += loss
    return total / iters

rng = np.random.default_rng(42)
t0 = time.time()
for step in range(steps + 1):
    xb, yb = get_batch(train_data, block_size, batch_size, rng)
    _, loss = model.forward(xb, yb)
    grads = model.backward()
    opt.step(model.P, grads)
    if step % 200 == 0:
        val = estimate_loss(model, val_data, rng)
        print(f"step {step:4d}  train {loss:.4f}  val {val:.4f}")
print("wall clock: %.1f s" % (time.time() - t0))

step    0  train 3.2026  val 2.9107
step  200  train 0.1950  val 0.1829
step  400  train 0.1548  val 0.1393
step  600  train 0.1446  val 0.1358
step  800  train 0.1460  val 0.1329
step 1000  train 0.1351  val 0.1389
step 1200  train 0.1112  val 0.1254
wall clock: 26.8 s

The loss falls from the random baseline of 3.20 to around 0.12 in about 27 seconds, and — the healthy sign — training and validation loss stay locked together the entire way (they even cross, with validation dipping slightly lower at some checkpoints because its batches happen to be easier). The repetitive corpus generalizes perfectly from the training split to the held-out split, so there’s no overfitting gap. This is a model that has genuinely learned its data.

Stage 3: See What It Learned#

The truest test is to let the model write. We do a simple greedy continuation: feed a seed prompt, take the single most likely next character, append it, and repeat. (Module 8 covers smarter sampling; greedy is enough to see whether anything was learned.)

def generate_greedy(model, prompt, n_new):
    ids = list(encode(prompt))
    for _ in range(n_new):
        context = np.array([ids[-block_size:]])
        logits, _ = model.forward(context)
        next_id = int(logits[0, -1].argmax())
        ids.append(next_id)
    return decode(ids)

print(repr(generate_greedy(model, "the ", 60)))

'the bay when the tide is low. the keeper counts the boats and wr'

From the seed "the ", the model produces the bay when the tide is low. the keeper counts the boats and wr — fully coherent Lantern Bay English, with correct spelling, spacing, sentence structure, and a period in the right place. Every word is real and the phrasing is lifted straight from the corpus’s grammar. That a 154,456-parameter model, built and trained by hand in NumPy in under half a minute, writes grammatical English at all is the whole point of the course. (Greedy decoding follows the single most likely path, so on this small corpus it tends to reproduce fluent stretches of the training text; the next module’s sampling methods trade a little of that fidelity for variety.)

Stage 4: Save and Reload the Trained Model#

A trained model is only useful if you can keep it. We serialize the parameter dictionary to a compressed .npz file (written to a temporary directory, not the repository), reload it into a fresh model, and confirm the reloaded model reproduces the original’s loss to the last digit.

import tempfile, os

path = os.path.join(tempfile.mkdtemp(), "mini_gpt.npz")
np.savez(path, **model.P)

reloaded = MiniGPT(vocab_size, block_size, n_embd, n_head, n_layer, seed=0)
weights = np.load(path)
for k in reloaded.P:
    reloaded.P[k] = weights[k]

check_rng = np.random.default_rng(7)
xb, yb = get_batch(train_data, block_size, batch_size, check_rng)
_, loss_original = model.forward(xb, yb)
_, loss_reloaded = reloaded.forward(xb, yb)
print("original loss:", round(float(loss_original), 4))
print("reloaded loss:", round(float(loss_reloaded), 4))
print("exact match:", abs(float(loss_original) - float(loss_reloaded)) < 1e-9)

original loss: 0.1216
reloaded loss: 0.1216
exact match: True

The reloaded model gives byte-for-byte the same loss, 0.1216, on the same batch — the save/load round-trip is lossless. This .npz file is your trained mini-GPT: a set of weights you can reload any time to generate text, which is exactly what Module 8 will do.

Practice Exercises#

Exercise 1: A smaller model#

Rerun with n_layer=2 and n_embd=32. How much does the final loss change, and how does the parameter count drop? Is the smaller model still enough for this corpus?

Exercise 2: Seed a different prompt#

Generate from the prompts "a small " and "the keeper ". Do the continuations stay on-theme? What does that say about what the model has captured?

Exercise 3: Confirm reproducibility#

Train two models with the same seed=42 and compare their step-0 and step-200 losses. Are they identical? Which single line makes the whole run reproducible?

Summary#

You ran a complete training session end to end: a 154,456-parameter mini-GPT trained in about 27 seconds of pure NumPy, its loss falling from the random baseline of 3.20 to about 0.12 with training and validation loss tracking together the whole way. A greedy continuation produced a coherent Lantern Bay sentence, and you saved the trained weights to disk and reloaded them to reproduce the model exactly. You now have a working, trained language model — built and trained entirely by hand.

Key Concepts#

• End-to-end run: configure → train while monitoring → generate → save.
• Healthy training: train and validation loss falling together, no widening gap.
• Greedy generation: repeatedly append the single most likely next token as a learning probe.
• Checkpointing: np.savez the parameter dict and reload it to reproduce the model losslessly.

Why This Matters#

This is the complete arc of training a neural language model, with nothing hidden: your own architecture, your own optimizer, your own training loop, producing a model that demonstrably learned real text and can be saved and reused. The workflow — pick a config, train while watching both losses, sanity-check by generating, checkpoint the result — is exactly the workflow behind models a million times larger. Having done it once, by hand, at a scale you can hold in your head, you understand what those larger runs are really doing.

Continue Building Your Skills#

You have a trained mini-GPT and a saved checkpoint, and you’ve seen it write with greedy decoding — which, as you noticed, is deterministic and prone to getting stuck. Module 8 is about generation itself: you’ll load this trained model and learn to sample from it well, with temperature to control randomness and top-k and top-p to keep samples coherent, turning the same weights into text that is varied instead of repetitive.