Train Your Own GPT

Lesson 4 ended with a promise: stop reading and train one. Today we keep it.

By the end of this lesson you will have taken the TinyGPT we built last time, pointed it at a text corpus small enough for a laptop, written the training loop with your own hands, and — this is the part I want you to actually watch, in real time, rather than take on faith — seen its output evolve from random characters, to word-shaped gibberish, to grammatical sentences in the voice of your corpus. Nothing I know of calibrates your intuitions about what these models are better than watching one condense out of randomness in front of you. It takes about an hour, most of which is the computer's time rather than yours.

What "training" actually means here

Let's assemble the pieces we already own. From Lessons 1–2: a neural network is a parameterised function, a loss measures how wrong it is on data, and gradient descent nudges every parameter slightly downhill on that loss, over and over. From Lesson 4: the transformer is such a function — it reads a token sequence and outputs, at every position, a probability distribution over the next token.

Training a GPT is nothing more than the marriage of the two:

The corpus

Frontier models train on trillions of tokens. We need something a laptop can chew through in minutes, which means roughly a few megabytes of plain text — and at this scale, the honest choice is character-level modelling: our vocabulary will be the set of distinct characters in the file (typically 60–100), not the 50,000-piece token vocabularies of Lesson 4. Every concept transfers unchanged; only the granularity differs.

Any plain-text file works — collected essays, public-domain novels, your own writing. Being who I am, I trained mine on a public-domain classic of medical literature: the 1918 edition of Gray's Anatomy, freely available from Project Gutenberg (about 5 MB of glorious Edwardian anatomical prose). Watching a neural network learn to hallucinate confident anatomy is an instructive experience for any clinician — more on that below. Save whatever you choose as corpus.txt next to your script.

The data pipeline

Ten lines. We build the character vocabulary, encode the whole corpus as integers, hold out the last 10% for validation, and write a function that serves random training snippets:

import torch

text = open("corpus.txt", encoding="utf-8").read()
chars = sorted(set(text))
vocab_size = len(chars)                      # typically 60–100 characters
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda t: "".join(itos[i] for i in t)

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]    # last 10% never trained on

ctx, batch_size = 256, 64

def get_batch(split):
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - ctx - 1, (batch_size,))
    x = torch.stack([d[i     : i+ctx    ] for i in ix])   # inputs
    y = torch.stack([d[i + 1 : i+ctx + 1] for i in ix])   # same text, shifted by 1
    return x.to(device), y.to(device)

Look hard at get_batch — it is Figure 1's step ① and the heart of self-supervision. The target y is literally the input x shifted one character. Because of the causal mask from Lesson 4, position t in the model's output only ever saw characters up to t, so predicting y[t] — the character at t+1 — is a fair exam, at all 256 positions of all 64 snippets simultaneously. One batch therefore contains 16,384 individual next-character exams.

The held-out 10% deserves emphasis. The model will never take a gradient step on it, which makes it our external validation cohort: performance there tells us the model is learning the language, not memorising the file. Every clinician who has watched a risk score validated only on its derivation cohort knows exactly why this split is non-negotiable.

The training loop

Bring in the TinyGPT from Lesson 4 unchanged, then:

import torch.nn.functional as F

device = ("mps" if torch.backends.mps.is_available()      # Apple Silicon
          else "cuda" if torch.cuda.is_available()        # NVIDIA
          else "cpu")

model = TinyGPT(vocab_size).to(device)     # ~800k parameters at default size
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)

@torch.no_grad()
def estimate_loss(split, iters=50):
    model.eval()
    losses = []
    for _ in range(iters):
        x, y = get_batch(split)
        logits = model(x)
        losses.append(F.cross_entropy(
            logits.view(-1, vocab_size), y.view(-1)).item())
    model.train()
    return sum(losses) / len(losses)

for step in range(5001):
    if step % 500 == 0:
        tr, va = estimate_loss("train"), estimate_loss("val")
        print(f"step {step:5d}   train {tr:.3f}   val {va:.3f}")

    x, y = get_batch("train")
    logits = model(x)                                          # ② forward
    loss = F.cross_entropy(logits.view(-1, vocab_size),        # ③ score
                           y.view(-1))
    opt.zero_grad()
    loss.backward()                                            # ④ gradients
    opt.step()                                                 # ⑤ nudge

That's the entire loop — compare it line by line against Figure 1. Two notes for the practitioners:

The optimizer is AdamW, not raw gradient descent. It's the same "step downhill" idea from Lesson 2, with two refinements earned by a decade of practice: each parameter gets an adaptive step size based on its recent gradient history, and a slight pull toward zero (weight decay) discourages overfitting. You could train with plain SGD; it would just take longer and land somewhere slightly worse.

The loss number is interpretable — use that. Cross-entropy here is the model's average surprise, in units of nats, at each true next character. Before training, with ~85 characters all equally likely, expect ln(85) ≈ 4.4. A model that has learned English character statistics lands around 2.0; a well-trained one on this corpus reaches 1.4–1.5 on validation. When your step-0 print shows ≈4.4, that's your first sanity check passing: the untrained model is exactly as ignorant as theory predicts.

Sampling: watching it think out loud

To generate text we run Lesson 4's Figure 1 literally — predict, sample, append, repeat:

@torch.no_grad()
def generate(model, prompt, max_new_tokens=400, temperature=1.0):
    idx = torch.tensor([encode(prompt)], device=device)
    model.eval()
    for _ in range(max_new_tokens):
        logits = model(idx[:, -ctx:])            # crop to context window
        logits = logits[:, -1, :] / temperature  # last position only
        probs = F.softmax(logits, dim=-1)
        nxt = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, nxt], dim=1)
    model.train()
    return decode(idx[0].tolist())

Now add print(generate(model, "The ")) inside the every-500-steps block, run the script, and watch. This is the part I promised. Roughly — your exact output will differ — here is what my Gray's Anatomy run produced:

Sit with that progression for a moment, because it compresses the entire deep-learning story into four snippets. At step 0: uniform randomness. By 500 the model has learned the statistics of English spelling — letter frequencies, vowel placement, word lengths — with no idea what a word means. By 2000 it produces real anatomical vocabulary in grammatical arrangements. By 5000 it writes fluent Edwardian anatomical prose that is anatomically wrong — that final sample is confident nonsense; the internal carotid does no such thing.

Reading the loss curves

While samples give you the qualitative story, the two loss numbers you're printing tell the quantitative one:

That divergence point is overfitting, and our 800k-parameter model on a 5 MB corpus will hit it eventually — the model is small relative to the data, so it holds off for a while, but it comes. The remedies are the classics: more data, a smaller model, more weight decay, or simply stopping when validation loss bottoms out. If this sounds exactly like the derivation-cohort-versus-validation-cohort discipline of clinical prediction models: yes. It is the same statistics wearing different clothes.

Practical notes for your run

On timing: with the default settings (4 layers, 128-wide, context 256), a 5,000-step run takes roughly 10–20 minutes on an Apple Silicon Mac (the script auto-selects the mps backend) or a modest NVIDIA GPU, and an hour or two on pure CPU. If CPU is all you have, drop ctx to 128 and the model to 3 layers — the qualitative arc from noise to sentences survives entirely intact.

On knobs worth turning once it works, in rough order of instructional value: temperature in generate() (0.5 is cautious and repetitive; 1.5 is creative and unhinged — an intuition worth having, since it's the same parameter you set in every LLM API); the learning rate (try 3e-3 and watch training destabilise — the divergence you'll see is Lesson 2's "overshooting the valley" made vivid); model width and depth (double both and watch validation loss improve — a miniature scaling law on your own hardware); and the corpus itself (swap in a different author; the model's entire personality follows the data — nothing else changed).

The complete script — data pipeline, Lesson 4's model, training loop, sampling, all assembled and commented — is here: train_tinygpt.py. Read it top to bottom before running it; it's ~130 lines and there is nothing in it you haven't now met.

What scaling changes — and what it doesn't

It's worth saying plainly what separates your afternoon's model from a frontier one, because the list is shorter than most people assume: more data (megabytes → tens of trillions of tokens), more parameters (800k → hundreds of billions), subword tokens instead of characters (Lesson 4), longer context, and a few architectural refinements per year of engineering. The training loop you just wrote — batch, forward, cross-entropy, backward, step — is, to a first approximation, what the big labs run, distributed across thousands of GPUs for months. The loss-versus-compute curves you sketched by doubling your model's width become, at industrial scale, the scaling laws that let labs forecast a model's loss before spending the money to train it.

And one thing scaling does not change: the model that comes out of this loop — at any size — is a raw next-token predictor with exactly the confident-fluency-without-fidelity property you bred in your own laptop run. It answers questions by continuing them, sometimes with more questions. Turning that raw material into something that reliably helps — that answers, follows instructions, declines gracefully, cites honestly — is a separate act of shaping, performed after pretraining.

What's next

That shaping is Lesson 6 — now live: post-training — how a base model becomes an assistant. Instruction tuning, feedback-based refinement, why the same weights can host such different behaviours, and what that pipeline means for the failure modes you'll meet in deployed clinical tools. The raw predictor you trained today is the "before" picture; next time we look at the "after," and at exactly what happens in between.

Until then: run the script. Change the corpus. Break the learning rate on purpose. The point of this lesson was never the 8 MB of weights you'll end up with — it's that "GPT" now names a process you have personally executed, end to end, rather than a product you consume.

— Neal