Building a Mini Language Model with PyTorch: Tutorial & Walkthrough

This is a build-it-yourself language model: about sixty lines of PyTorch that learns to write Shakespeare one character at a time. No pretrained weights, no library beyond torch, runs on a laptop in ten minutes.

I wrote the first version of this post in 2024 and it had a flaw I want to fix, because it is the flaw most tutorials have. It trained on the string “Hello world. Hello again.” and then stopped, before anything interesting could happen. You got a model that compiled. You did not get to see it learn, which is the only part worth watching.

So this time everything is a real run, and the numbers below are from that run rather than from memory.

The setup

1.1 MB of Shakespeare, 65 distinct characters, 90/10 train/val split. Characters rather than words because it keeps the vocabulary at 65 instead of 30,000 and lets the whole thing stay small enough to read.

text = open("input.txt").read()
chars = sorted(set(text))                       # 65 of them
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]


def get_batch(split):
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - BLOCK - 1, (BATCH,))
    x = torch.stack([d[i:i + BLOCK] for i in ix])
    y = torch.stack([d[i + 1:i + BLOCK + 1] for i in ix])   # shifted by one
    return x, y

That one-character shift in y is the entire supervised signal. The label for every position is the character that came next, which is why this needs no annotation and why it scales to any text you can find.

The model

Three pieces: turn characters into vectors, mix information across the sequence, project back to a score per character.

class RNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(VOCAB, EMBED)
        self.rnn = nn.RNN(EMBED, EMBED, num_layers=2, batch_first=True)
        self.head = nn.Linear(EMBED, VOCAB)

    def forward(self, idx):
        x = self.embed(idx)        # (B, T) -> (B, T, EMBED)
        out, _ = self.rnn(x)       # mix across time
        return self.head(out)      # (B, T, VOCAB) logits

0.17M parameters. The middle line is the only architectural claim in the whole file, and it is the line I am going to replace later.

for step in range(STEPS + 1):
    x, y = get_batch("train")
    logits = model(x)
    loss = F.cross_entropy(logits.reshape(-1, VOCAB), y.reshape(-1))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

The gradient clipping is not decoration. Take it out of the RNN run and you will eventually meet the exploding-gradient problem in person: a loss that goes to nan and stays there.

Watching it learn

The useful thing about a character model is that you can read its output at any point and know exactly how far along it is. Sampling from the same seed at three moments:

Step 0: loss 4.19. This is ln(65), the loss of guessing uniformly. The model knows nothing, and the sample is what nothing looks like:

khR$snpE;M3bZoPcxfEeCnIhjf.LV
dbT??t3BlO-qIwRFc:.m'i,B-oduK gjgS,s-jBPPWBhAVJJdEvyr

Step 500: loss 1.96. In eighty seconds it has learned that text is made of words, that words are separated by spaces, that lines end, and roughly which letters follow which. It has learned English orthography without being told English exists:

DI shile, with I finiss
prentle earver his more flaints, we ding to poinmed; an more of-ere'd blood
thy his ways Do$k some, hecours; stame foak if compone,

Step 9000: loss 1.55. Now it has the format. Character names in caps, colon, newline, speech. Grammar that holds for a clause or two before dissolving:

Smam, hang's shalt have too dear up.

TUKENCE!' Troth I have been in seem tongue by and out of mine.

ISABELLA:
Poor better did I, my majession; and serve so moningham;
And we do attending of you gone and fain,

“Poor better did I, my majession” is nonsense, but it is nonsense with the right shape, and the model got there from a table of character frequencies in ten minutes. That progression (noise, then orthography, then structure) is the thing the original version of this post never let you see.

Then I swapped the recurrence for attention

The obvious question in 2026 is why build an RNN at all. So I built the other one too: same data, same batch size, same context length, same optimiser, same seed. Only the middle line changes.

class Block(nn.Module):
    def forward(self, x, mask):
        h = self.ln1(x)
        a, _ = self.attn(h, h, h, attn_mask=mask)   # every position sees every
        x = x + a                                   # earlier position directly
        return x + self.mlp(self.ln2(x))

Instead of squeezing the past through a fixed-size hidden state one step at a time, every position reads every earlier position directly. 0.94M parameters, five and a half times the RNN.

I expected it to win comfortably. Here is what actually happened.

	RNN	Transformer
Parameters	0.17M	0.94M
Val loss @ 3000	1.6489	1.6746
Val loss @ 9000	1.5547	1.5580
Train loss @ 9000	1.3372	1.2632
Train–val gap	0.218	0.295
Wall clock @ 9000	597s	406s

They tie. The transformer is behind at 3000 steps, catches up, and finishes 0.003 nats worse, which is noise. Five and a half times the parameters bought nothing on the validation set.

It did buy something on the training set. Train loss went to 1.2632 against the RNN’s 1.3372, so the gap between the two curves is wider. The bigger model is not learning more about English. It is memorising more of this particular Shakespeare.

At 1.1 MB, the architecture is not the bottleneck. The data is. Both models have already extracted about as much as this corpus contains, and the larger one is spending its extra capacity on remembering rather than generalising.

This is the opposite of the advice the earlier version of this post gave. It had a section suggesting you add layers and widen embeddings if the model underfits. On a dataset this size that advice will move your train loss and leave your val loss alone, and you will feel like you are making progress for an afternoon.

It is also, I think, the honest version of the scaling story. Transformers did not win because attention is magic on a megabyte. They won because attention parallelises across the sequence (the transformer finished the same 9000 steps in 406 seconds against the RNN’s 597, despite being 5.5× larger), and that throughput is what let people train on far more than a megabyte. The advantage is real at 10 GB. It is invisible at 1 MB.

Things worth trying, in order

If you run this yourself, these are the changes that will actually teach you something, roughly in order of information gained per minute spent:

• Remove the gradient clipping from the RNN run and watch where it diverges. Five minutes, and you will never forget what exploding gradients look like.
• Print the train–val gap every eval rather than just the losses. It is the number that tells you whether to get more data or a bigger model, and almost no tutorial plots it.
• Train on ten times more text. Concatenate more Gutenberg. This is where the transformer starts to pull away, and seeing that happen is more convincing than being told.
• Change the sampling temperature. Divide the logits by 0.5 and by 1.5 before the softmax. Same model, and the difference between confident-and-repetitive and creative-and-incoherent is one division.
• Swap characters for a subword tokeniser. The vocabulary goes from 65 to a few thousand, the sequences get shorter, and the same compute buys much more context.

Caveats

• One seed, one run per architecture. The 0.003 difference in final val loss is well inside seed variance, which is exactly why I am calling it a tie rather than an RNN win.
• Neither model was tuned. Same learning rate for both, chosen because it was stable, not because it was best for either. A tuned transformer would likely beat a tuned RNN by more than this; the point stands that at this data scale the margin is small and the overfitting gap is the thing to watch.
• Run on an M4 Pro through the MPS backend. The wall-clock numbers are hardware-specific; the ratio between them is the transferable part.

The full script is one file: two model classes, a batcher, a training loop, and a sampler. If you want the version of this exercise that goes further, Karpathy’s nanoGPT and his build-GPT video are the canonical next step, and this post is unashamedly downstream of them.