Build and train a small transformer from scratch
By the end: You have written, trained and sampled from a character-level transformer, and can explain every line.
- Data as a stream of token ids
- Batches of context and next-token targets
- A transformer block in twenty lines
- The training loop with validation
- Sampling with temperature
New to pytorch? PyTorch is the Python library almost everyone uses to build and train neural networks. It hands you two things you would otherwise have to write: arrays of numbers that can live on a GPU instead of the main processor, and automatic differentiation — it watches the arithmetic you perform and works out on its own how every weight should move to reduce the error. Layers, optimisers and data loading come with it. Why it matters
You know roughly how a transformer works; this is writing one, in about eighty lines, and training it on a book until it produces text. Embeddings, a causal attention block, residual adds, a training loop with validation, and sampling with a temperature: the same structure as the models you call over an API, with the scale taken out. Writing it yourself is what turns the architecture diagram into something you can debug, and it is the piece a research interview is most likely to ask you to reproduce live.
Piece In this file In a production LLM
tokeniser one integer per char BPE, 100k+ merged subwords
context length T 128 128k to 1M
width D 256 4,000 to 16,000
heads H 4 32 to 128
layers L 4 32 to 120
parameters 3.2 million billions
position information learned embedding RoPE, applied inside attention
normalisation LayerNorm, pre-norm RMSNorm, pre-norm
attention dense causal grouped-query, with a KV cache
training data one book trillions of tokensData#
Use any plain text file of a few megabytes, such as a public domain book. Each character becomes an integer. The task is to predict the next character from the ones before it.
Characters rather than subwords because it removes a whole component: no tokeniser to train, no vocabulary file, and a vocabulary of about 65 that you can print. The cost is that the model spends its capacity learning spelling before it learns anything else, which is why the samples are gibberish for the first thousand steps and then abruptly contain real words.
import torch, torch.nn as nn, torch.nn.functional as F
text = open("book.txt", encoding="utf-8").read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
split = int(0.9 * len(data))
train, val = data[:split], data[split:]
V, T, D, H, L = len(chars), 128, 256, 4, 4 # vocab, context, width, heads, layers
def batch(src, n=32):
ix = torch.randint(len(src) - T - 1, (n,))
x = torch.stack([src[i:i + T] for i in ix]) # [n, T]
y = torch.stack([src[i + 1:i + T + 1] for i in ix]) # the same, shifted by one
return x.cuda(), y.cuda()The y line is the entire supervision signal and it is worth staring at. The target is the input shifted by one, so position 0 must predict what sits at position 1, position 1 predicts position 2, and so on. One sequence of 128 characters therefore provides 128 separate training examples, all computed in a single forward pass. That is why language models are efficient to train and why the causal mask is not optional.
The model#
class Block(nn.Module):
def __init__(self):
super().__init__()
self.ln1, self.ln2 = nn.LayerNorm(D), nn.LayerNorm(D)
self.qkv = nn.Linear(D, 3 * D)
self.proj = nn.Linear(D, D)
self.mlp = nn.Sequential(nn.Linear(D, 4 * D), nn.GELU(), nn.Linear(4 * D, D))
def forward(self, x): # [B, T, D]
B, T_, _ = x.shape
q, k, v = self.qkv(self.ln1(x)).split(D, dim=-1)
q, k, v = (t.view(B, T_, H, D // H).transpose(1, 2) for t in (q, k, v))
a = F.scaled_dot_product_attention(q, k, v, is_causal=True) # [B, H, T, D/H]
x = x + self.proj(a.transpose(1, 2).reshape(B, T_, D)) # attention, added
return x + self.mlp(self.ln2(x)) # MLP, added
class TinyGPT(nn.Module):
def __init__(self):
super().__init__()
self.tok, self.pos = nn.Embedding(V, D), nn.Embedding(T, D)
self.blocks = nn.Sequential(*[Block() for _ in range(L)])
self.ln, self.head = nn.LayerNorm(D), nn.Linear(D, V)
def forward(self, idx): # [B, T] of ids
pos = torch.arange(idx.shape[1], device=idx.device)
x = self.tok(idx) + self.pos(pos)
return self.head(self.ln(self.blocks(x))) # [B, T, V] logitsidx [32, 128] token ids
tok(idx) + pos(arange) [32, 128, 256] embeddings, added not concatenated
per block:
ln1(x) [32, 128, 256]
qkv(...) [32, 128, 768] one matmul for all three
.split(D, dim=-1) 3 x [32, 128, 256]
.view(B, T, H, D//H) [32, 128, 4, 64] split the width across heads
.transpose(1, 2) [32, 4, 128, 64] heads become a batch dimension
attention [32, 4, 128, 64] internally [32, 4, 128, 128] scores
.transpose(1,2).reshape(...) [32, 128, 256] heads rejoined
x + proj(...) [32, 128, 256] residual
x + mlp(ln2(x)) [32, 128, 256] residual
head(ln(x)) [32, 128, 65] a distribution per positionTwo design choices in that walkthrough are worth naming because they recur everywhere. Q, K and V come from one Linear of width 3D rather than three of width D, because a single large matrix multiply is faster than three small ones and the split afterwards is free. And multi-head attention is not several attention modules; it is one, with the width reshaped so the head index behaves like an extra batch dimension. Nothing in the attention code knows there are heads.
Map it to what you know
Every piece here appears in large language models: token and position embeddings, pre-norm blocks with a residual add after attention and after the MLP, a causal mask, a final norm and a projection to vocabulary logits. Scale changes; the structure does not.
token embedding 65 x 256 16,640
position embedding 128 x 256 32,768
per block:
2 LayerNorms 2 x 2 x 256 1,024
qkv 256 x 768 + 768 197,376
proj 256 x 256 + 256 65,792
mlp 256x1024 + 1024 + 1024x256 + 256 525,568
per block 789,760
x 4 blocks 3,159,040
final norm + head 512 + 256 x 65 + 65 17,217
TOTAL 3,225,665Memory, separately, is dominated by something that has no parameters at all. The attention scores are [batch, heads, T, T], which here is 32 by 4 by 128 by 128, or 2.1 million numbers per layer held for the backward pass. Double the context length and that term quadruples while the parameter count does not move. That single fact is why context length is expensive, why scaled_dot_product_attention is used here instead of writing the softmax by hand, and why FlashAttention exists.
Train it#
model = TinyGPT().cuda()
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
for step in range(5001):
x, y = batch(train)
loss = F.cross_entropy(model(x).view(-1, V), y.view(-1))
opt.zero_grad(); loss.backward(); opt.step()
if step % 500 == 0:
model.eval()
with torch.no_grad():
vx, vy = batch(val)
vloss = F.cross_entropy(model(vx).view(-1, V), vy.view(-1))
model.train()
print(f"step {step:5} train {loss.item():.3f} val {vloss.item():.3f}")At step 0 the loss should be close to the natural log of the vocabulary size, about 4.2 for 65 characters, because an untrained model guesses uniformly. If it starts far from that, something is wrong before training even begins.
Run on a smaller corpus with a 28-character vocabulary, on CPU:
loss at init 3.4867 ln(28) = 3.3322 close enough: the model is uninformed
logits shape (32, 128, 28) one distribution per position
step 0 train 3.494 val 3.128
step 100 train 0.171 val 0.192
step 200 train 0.021 val 0.023
step 300 train 0.014 val 0.016
sample: 'the n liquor jugs. pack my box with five doze dozen liquor jugs. pack my...'Running it on something trivial first is a habit worth keeping. A model that cannot memorise two sentences has a bug, and you find that out in ninety seconds rather than after an overnight run on a real book. It is the same argument as overfitting a single batch, applied to the whole pipeline including sampling.
On a real book expect a different shape: validation loss around 1.5 after a few thousand steps, samples that are pronounceable nonsense at step 500, recognisable English words by 2,000, and something with the rhythm of the source and no meaning by 5,000. Validation loss rising while training loss falls means the model is large for the corpus, which at 3.2 million parameters on a single book it genuinely is.
Sample from it#
@torch.no_grad()
def sample(start, n=400, temperature=0.8):
model.eval()
idx = torch.tensor([[stoi[c] for c in start]], device="cuda")
for _ in range(n):
logits = model(idx[:, -T:])[:, -1] / temperature # last position only
nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return "".join(chars[i] for i in idx[0].tolist())
print(sample("The "))Four details in that function are the whole of inference. The model is run on idx[:, -T:] because it has only T position embeddings and cannot be given more; a longer prompt must be truncated. Only the last position's logits are used, because the others are predicting characters you already have. multinomial samples from the distribution rather than taking the argmax, which is what stops the output repeating a single phrase forever. And each new token is appended and the whole thing re-run, which is why generation is sequential and slow.
That last point is also the one real inefficiency here, and naming it is worth marks. Every step recomputes the keys and values for all the preceding positions, having computed exactly the same numbers on the previous step. Caching them is the KV cache, and it turns generation from quadratic in the output length into linear. This implementation has none, which is fine for 400 characters and hopeless for a chat model.
temperature effect when
0.0 (argmax) fully deterministic, repeats extraction, classification, code completion
0.7 to 0.8 coherent with some variety the usual default
1.0 the model's raw distribution when you want the true samples
> 1.2 increasingly incoherent creative tasks, and not much elseTemperature 0 is not a temperature
Dividing by zero is undefined, so every implementation special-cases it to mean argmax. Handle it with an if rather than passing 1e-8 and hoping. torch.softmax survives 1e-8 because it subtracts the row maximum internally, but a hand-written exp(logits / t) / sum at that temperature overflows to NaN across the whole row, which is the failure the max subtraction exists to prevent.
Where people go wrong
Copying a reference implementation and running it. The learning is in writing it yourself and fixing the shape errors you hit.
If this comes up in an interview
Walk me through a minimal GPT-style model.
Token ids go through token and position embeddings and are added. A stack of blocks follows, each with layer norm, causal multi-head self-attention added back to the residual, then layer norm and an MLP added back. A final norm and a linear layer produce logits over the vocabulary. Training minimises cross-entropy between logits and the next token at every position.
Then they push further
Why is the loss at initialisation a useful check?
An untrained model should predict roughly uniformly, so its loss should be close to the log of the vocabulary size. Much higher means a bad initialisation or a bug, such as targets not shifted correctly.
Why is the causal mask needed during training?
Training predicts every position at once. Without the mask, each position could attend to the tokens it is supposed to predict, and the model would learn to copy rather than to predict.
What gets marked up. Strong candidates can write this from memory and explain residual adds, pre-norm and the mask. Many applicants for research roles are asked to implement part of it live.
Build this: Type the model out yourself, train it for 5,000 steps on a public domain book, and sample. Then double the context length, then the layers, and record validation loss and training time for each change.
Watch it explained
This is one section of PyTorch, 30 hours across 6 sections. Read it in full to get the whole topic in order.