Hands-On: Build a Tiny GPT (Transformers Demystified in Code)
Introduction
You've learned every piece of the Transformer — tokens, embeddings, positional encoding, self-attention, multi-head, blocks. Now we assemble them into a working GPT in about 100 lines of PyTorch, train it, and watch it generate text. This is the moment the magic fully dissolves: you'll see there's no magic at all, just the pieces you already understand, wired together.
This is a learning toy (character-level, tiny), in the spirit of Andrej Karpathy's nanoGPT. It is the same architecture as a frontier model — just unimaginably smaller. Get this running and you can honestly say you've built a GPT from scratch.
You'll build:
- A char-level tokenizer and data batches
- The model: embeddings + positional + causal multi-head attention + feed-forward blocks
- A training loop (next-character prediction)
- A generator (the predict → sample → append loop)
Prereq: pip install torch, and any plain-text file saved as input.txt (the classic choice is "tiny Shakespeare").
The Plan (You Already Know All of This)
Before any code, notice that the whole build is just this section's lessons, in order:
| Step | Concept (lesson) |
|---|---|
| Char tokenizer | Tokenization |
| Token + position embeddings | Embeddings + Positional Encoding |
| Causal self-attention head | Self-Attention (causal) |
| Several heads in parallel | Multi-Head Attention |
| Block = attention + feed-forward + add&norm | The Transformer Architecture |
| Stack N blocks → linear → softmax | Refinement stack → Next-Token Prediction |
| Train on next-char, then generate | Next-Token Prediction (autoregressive) |
Nothing here is new. We're just writing it down.
Step 1 — Data & a Char-Level Tokenizer
To keep it dead simple, our "tokens" are single characters (no BPE). We map each unique character to an integer and back, then make batches where the target of each position is simply the next character.
import torch
torch.manual_seed(1337)
text = open('input.txt').read() # any plain text (e.g. tiny Shakespeare)
chars = sorted(set(text)); vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)} # char -> int
itos = {i: c for c, i in stoi.items()} # int -> char
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data)); train_data, val_data = data[:n], data[n:]
block_size = 64 # context length (how many chars the model sees)
batch_size = 32
def get_batch(split):
d = train_data if split == 'train' else val_data
ix = torch.randint(len(d) - block_size, (batch_size,))
x = torch.stack([d[i:i+block_size] for i in ix])
y = torch.stack([d[i+1:i+block_size+1] for i in ix]) # targets = next char
return x, yStep 2 — The Model (the whole Transformer)
Here's the heart of it. Read it against your mental model — every class is a lesson you've done. First, one self-attention head (the softmax(Q·Kᵀ/√d)·V you saw, plus the causal mask), and multi-head (run several, concatenate):
import torch.nn as nn
from torch.nn import functional as F
n_embd, n_head, n_layer = 128, 4, 4
head_size = n_embd // n_head # the dimension is SPLIT across heads
class Head(nn.Module): # one self-attention head
def __init__(self):
super().__init__()
self.key = nn.Linear(n_embd, head_size, bias=False) # learned K
self.query = nn.Linear(n_embd, head_size, bias=False) # learned Q
self.value = nn.Linear(n_embd, head_size, bias=False) # learned V
self.register_buffer('mask', torch.tril(torch.ones(block_size, block_size)))
def forward(self, x):
B, T, C = x.shape
k, q, v = self.key(x), self.query(x), self.value(x)
att = q @ k.transpose(-2, -1) * head_size**-0.5 # relevance scores
att = att.masked_fill(self.mask[:T,:T] == 0, float('-inf')) # CAUSAL: look back only
att = F.softmax(att, dim=-1) # -> attention weights
return att @ v # weighted blend of values
class MultiHead(nn.Module): # several heads in parallel = specialists
def __init__(self):
super().__init__()
self.heads = nn.ModuleList([Head() for _ in range(n_head)])
self.proj = nn.Linear(n_embd, n_embd)
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1) # concatenate heads
return self.proj(out) # project backNow the feed-forward net, a block (attention + feed-forward, each wrapped in Add & Norm), and the full GPT — token+position embeddings, a stack of blocks, and a final linear head to next-char scores:
class FeedForward(nn.Module): # the 'think' step (a small MLP)
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4*n_embd), nn.ReLU(), nn.Linear(4*n_embd, n_embd))
def forward(self, x): return self.net(x)
class Block(nn.Module): # attention -> feed-forward, with residual + layernorm
def __init__(self):
super().__init__()
self.sa, self.ff = MultiHead(), FeedForward()
self.ln1, self.ln2 = nn.LayerNorm(n_embd), nn.LayerNorm(n_embd)
def forward(self, x):
x = x + self.sa(self.ln1(x)) # residual + attention (Add & Norm)
x = x + self.ff(self.ln2(x)) # residual + feed-forward (Add & Norm)
return x
class TinyGPT(nn.Module):
def __init__(self):
super().__init__()
self.tok_emb = nn.Embedding(vocab_size, n_embd) # token meaning
self.pos_emb = nn.Embedding(block_size, n_embd) # token position
self.blocks = nn.Sequential(*[Block() for _ in range(n_layer)])
self.ln_f = nn.LayerNorm(n_embd)
self.head = nn.Linear(n_embd, vocab_size) # -> a score per char
def forward(self, idx, targets=None):
B, T = idx.shape
x = self.tok_emb(idx) + self.pos_emb(torch.arange(T)) # meaning + position
x = self.ln_f(self.blocks(x)) # the refinement stack
logits = self.head(x)
loss = None
if targets is not None:
loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
return logits, loss
@torch.no_grad()
def generate(self, idx, max_new=300):
for _ in range(max_new):
logits, _ = self(idx[:, -block_size:]) # predict
probs = F.softmax(logits[:, -1, :], dim=-1) # distribution over next char
nxt = torch.multinomial(probs, 1) # sample one
idx = torch.cat([idx, nxt], dim=1) # append -> repeat
return idxThat's the entire model. Look how it lines up with the architecture diagram: tok_emb + pos_emb (meaning + position) → blocks (N× attention + feed-forward) → ln_f + head (final linear) → softmax in generate. You've seen every line conceptually — now it's code.
Step 3 — Train It (next-character prediction)
Training is the self-supervised loop from the foundations lesson: predict the next char, compare to the real one, nudge the weights, repeat. Watch the loss drop — that's the model learning.
model = TinyGPT()
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
for step in range(5000):
xb, yb = get_batch('train')
_, loss = model(xb, yb)
opt.zero_grad(); loss.backward(); opt.step() # backprop + update
if step % 500 == 0:
print(f"step {step:4d} | loss {loss.item():.3f}")
# step 0 | loss 4.21 (random — ~ln(vocab_size))
# step 500 | loss 2.49
# step 5000 | loss 1.61 (it has learned the statistics of the text)Step 4 — Generate Text
Now sample from it — the autoregressive loop you've seen since lesson 1: predict a distribution, sample a character, append, repeat.
start = torch.zeros((1, 1), dtype=torch.long) # begin from a single token
print(decode(model.generate(start, max_new=500)[0].tolist()))
Out of random characters, the model learns spelling, spacing, names, and punctuation — purely from predicting the next character. No grammar rules were programmed; it absorbed them from the text. That's the whole thesis of this section, demonstrated in code you wrote.
What You Just Did
Take a second — you just implemented, from scratch:
- Tokenization (char-level) and batching
- Token + positional embeddings
- Causal multi-head self-attention (
softmax(Q·Kᵀ/√d)·V+ a look-back mask) - Feed-forward networks, residual connections, and layer norm
- A stack of Transformer blocks
- Self-supervised training (next-token loss + backprop)
- Autoregressive generation (predict → sample → append)
That is the entire architecture of a modern LLM. You're not looking at GPT from the outside anymore — you've built one.
So What's Different in a Real LLM?
If the architecture is the same, why is GPT-4 so much more capable than our toy? Scale and polish — not a different machine:
- Size. Billions of parameters and dozens-to-hundreds of layers (we used 4 layers, 128 dims).
- Data. Trillions of tokens of high-quality text (we used one small file).
- Tokenizer. Subword BPE, not characters (far more efficient).
- Positional encoding. RoPE instead of simple learned positions.
- Engineering. Distributed training, careful optimization, huge compute (reproducing GPT-2 today costs only a few hundred dollars of compute; frontier models cost millions).
- Post-training. SFT + RLHF/DPO to make it follow instructions and be helpful (the entire next sections).
The profound takeaway: the architecture you just built is the architecture. Everything else is scale, better data, and post-training — all of which we cover next.
See It: Train Your Tiny GPT
Before you run the real code, run the toy. Click train and watch the loss drop as the generated text sharpens from random noise → word-shaped → coherent Shakespeare. Then break it: shrink the model or the context length and feel why scale matters.
Nothing was programmed to know English — trained only to predict the next character, the model absorbed spelling, names, and grammar from the text. A frontier LLM is this same machine, just with far more size, data, and training — exactly what the next sections cover.
🧪 Try It Yourself
Your turn — break it to understand it. In the tiny-GPT code, try one change and predict the effect before running:
- Shrink the context length (block size) to 8 → what can it no longer learn? (Long-range structure.)
- Set embedding size very small → output gets more garbled. Why? (Less room to represent meaning.)
- Train for 10× more steps → loss drops, samples sharpen.
Seeing your edits move the loss is the fastest way to make 'it's just next-token prediction' click.
Key Takeaways
- A working GPT is ~100 lines: char tokenizer → token+position embeddings → N causal multi-head attention blocks → linear + softmax → train on next-char → generate.
- Every line maps to a concept from this section — there is no hidden magic.
- Trained only on next-character prediction, it learns spelling, grammar, and structure on its own.
- A frontier LLM is the same architecture, massively scaled (more params/data/compute) plus better tokenization/position and post-training — not a different design.
You've completed the foundations of how AI works. Next, we zoom back out: Foundation Models Demystified — pretraining at internet scale, scaling laws, and the model families (GPT, Claude, Llama, …) you'll actually build with.