Build an LLMMiddleLesson 417 min read

Step 5 · A complete tiny GPT

Every piece assembled into a model you can actually train tonight. Under 100 lines, and structurally identical to a frontier model.

🧪 Lab · 45 min 📝 Lesson test

Lesson in motion

In 60 seconds

Step 5 · A complete tiny GPT

Every piece assembled into a model you can actually train tonight. Under 100 lines, and structurally identical to a frontier model.

1/5

In simple words

Now we glue the parts together: the number-lookup, forty thinking blocks, and a final layer that guesses the next chunk.

The whole modelpython

import torch, torch.nn as nn, torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class Config:
    vocab_size: int = 8192
    n_layer: int = 6
    n_head: int = 6
    d_model: int = 384
    block_size: int = 256      # max context length
    dropout: float = 0.1

class TinyGPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos = nn.Embedding(cfg.block_size, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList(
            Block(cfg.d_model, cfg.n_head) for _ in range(cfg.n_layer))
        self.norm = RMSNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.tok.weight        # weight tying
        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok(idx) + self.pos(pos))
        for b in self.blocks:
            x = b(x)
        logits = self.head(self.norm(x))          # (B, T, vocab)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new=100, temperature=1.0, top_k=50):
        for _ in range(max_new):
            idx_cond = idx[:, -self.cfg.block_size:]      # crop to context
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature       # last position only
            if top_k:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = -float("inf")
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, nxt), dim=1)
        return idx

Read the two important lines again

1
The targets are the inputs, shifted by one
targets = idx[1:]. That is the entire training objective. Predict the next token, everywhere in the sequence at once. There is no other supervision.
2
Generation is a loop over that same forward pass
Take the logits at the last position, sample one token, append it, run again. Everything an LLM ever produces comes out of this loop.

Count your parameters before you train

Where the parameters livepython

def count(cfg):
    emb  = cfg.vocab_size * cfg.d_model                 # tied with output head
    attn = 4 * cfg.d_model * cfg.d_model                # q, k, v, proj
    ff   = 3 * cfg.d_model * int(8 * cfg.d_model / 3)   # gate, up, down
    per_layer = attn + ff
    total = emb + cfg.n_layer * per_layer
    print(f"embedding  {emb/1e6:.1f}M")
    print(f"per layer  {per_layer/1e6:.1f}M  x {cfg.n_layer}")
    print(f"TOTAL      {total/1e6:.1f}M")
    return total

count(Config())     # ~ 14M parameters

Watch out

For a small model the embedding table is often the single largest component. If your 14M-parameter model has 3M parameters of embedding, shrinking the vocabulary is the cheapest capacity win available.

What to expect when you run it

Training time	Loss	What it produces
30 seconds	~5.5	Random characters
2 minutes	~3.0	Word-shaped nonsense with real spacing
10 minutes	~2.0	Real words, broken grammar
1 hour	~1.5	Fluent-sounding sentences that mean nothing
Overnight, small dataset	~1.2	Memorising — check for overfitting

Do this

That progression is the single most valuable thing in this track to see with your own eyes. Watching gibberish turn into grammar over twenty minutes permanently changes how you think about these systems. Nothing else in the guide replaces it.

Watch and read more

**Let's reproduce GPT-2 (124M)**Andrej Karpathy · 4 hr · video

**Let's build GPT: from scratch**Andrej Karpathy · video

nanoGPTkarpathy/nanoGPT · reading

Lab

A language model you trained yourself, generating real words.

~45 min

The problem

Train the TinyGPT from Module 41 on a few MB of text. Sample from it every 200 steps and keep every sample. Watch gibberish become words become grammar. Then plot loss against sample quality.

Starter codepython

if step % 200 == 0:
    ctx = torch.zeros((1,1), dtype=torch.long, device=device)
    print(f"--- step {step} loss {loss.item():.3f} ---")
    print(decode(model.generate(ctx, max_new=120)[0].tolist()))

You are done when

Loss fell below 2.0
You have samples from at least six points in training
You can point at the step where words appeared

Hard questions

Try to answer before you reveal. If you can answer these, you understood the lesson.

Q1Your training loss keeps falling but samples stop improving. Give the two most likely causes and how to tell them apart.Reveal

Memorisation: the model is fitting your small corpus rather than learning language — check by comparing training and validation loss; a widening gap confirms it. Or you are sampling badly: greedy decoding on a small model produces repetitive text regardless of quality — check by varying temperature and top-k. The first needs more or better data; the second needs no training change at all, which is why you check it first.

Please sign in to continue.

Questions people ask

What dataset should I use?

Start with a single text file of a few megabytes — a public-domain book collection works well. Small enough to iterate in minutes, large enough that the model cannot simply memorise it. Then move to something like TinyStories or a small web-text sample.

Why is my loss stuck around 10?

Almost always a data or shape bug rather than a model bug. Check that targets are shifted by exactly one, that you are not accidentally feeding padding as targets, and that your token ids are within the vocabulary range.

Can I run this on a laptop?

Yes, at this size. Apple Silicon via the MPS backend or plain CPU will train the 14M configuration on a small file in tens of minutes. Reduce block_size if memory is tight.

How is this different from GPT-2?

Scale, and a few refinements. GPT-2 small is 124M parameters, 12 layers, d_model 768, context 1024, trained on 40 GB of text. This is the same architecture with the numbers turned down and RMSNorm and SwiGLU swapped in.

Lesson test

5 questions. Get 3 right (60%) to pass and complete this lesson.

Step 5 · A complete tiny GPT

Step 5 · A complete tiny GPT

Read the two important lines again

The targets are the inputs, shifted by one

Generation is a loop over that same forward pass

Count your parameters before you train

What to expect when you run it

Watch and read more

The problem

You are done when

Hard questions

Questions people ask