Module 10 — Implementing a complete Transformer block
Every module so far added a piece. This one puts them together into a single file that trains on a laptop and produces readable attention maps. The task is deliberately small — translate a spelled-out date into ISO format — but the machinery is exactly the one that powers real models.
The task, restated
Input: a date in a mixed English / numeric format, up to twenty characters long.
"3 March 2026"→"2026-03-03""12 Jan 2024"→"2024-01-12""5 September 2020"→"2020-09-05"
Character-level vocabulary of about 40 symbols. Training pairs are generated on the fly, so the dataset is effectively infinite and free. The point is not to translate dates — you would write ten lines of regex for that. The point is that the correct answer is visible in the input at specific positions, so a working Transformer produces attention maps we can read.
Building the vocabulary and dataset
import random
import torch
from torch.utils.data import Dataset, DataLoader
MONTHS = ["Jan", "Feb", "Mar", "Apr", "May", "Jun",
"Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
CHARS = list(" 0123456789-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
PAD, SOS, EOS = "<p>", "<s>", "<e>"
VOCAB = [PAD, SOS, EOS] + CHARS
STOI = {c: i for i, c in enumerate(VOCAB)}
ITOS = {i: c for c, i in STOI.items()}
def encode(s):
return [STOI[c] for c in s]
def sample_pair():
y, m, d = random.randint(2000, 2029), random.randint(1, 12), random.randint(1, 28)
src = f"{d} {MONTHS[m - 1]} {y}"
tgt = f"{y:04d}-{m:02d}-{d:02d}"
return src, tgt
class DateDataset(Dataset):
def __init__(self, n=20000):
self.pairs = [sample_pair() for _ in range(n)]
def __len__(self):
return len(self.pairs)
def __getitem__(self, i):
s, t = self.pairs[i]
src = torch.tensor(encode(s.ljust(20)))
tgt = torch.tensor(encode(SOS + t + EOS))
return src, tgt
ljust(20) pads the source to fixed length; the target keeps <s> and <e> markers so the decoder learns where to start and stop.
The full Transformer, using the modules we built
The file below imports the blocks from earlier modules and adds only the top-level plumbing. Every part has been justified.
import torch
import torch.nn as nn
from embedding import TokenAndPositional
from mha import MultiHeadAttention
from encoder_layer import FeedForward, SubLayerConnection
class EncoderLayer(nn.Module):
def __init__(self, d, h, drop):
super().__init__()
self.mha = MultiHeadAttention(d, h)
self.ff = FeedForward(d, dropout=drop)
self.s1 = SubLayerConnection(d, drop)
self.s2 = SubLayerConnection(d, drop)
def forward(self, x, mask):
x = self.s1(x, lambda y: self.mha(y, mask)[0])
return self.s2(x, self.ff)
class DecoderLayer(nn.Module):
def __init__(self, d, h, drop):
super().__init__()
self.self_a = MultiHeadAttention(d, h)
self.cross_a = MultiHeadAttention(d, h) # accepts separate q, kv
self.ff = FeedForward(d, dropout=drop)
self.s1 = SubLayerConnection(d, drop)
self.s2 = SubLayerConnection(d, drop)
self.s3 = SubLayerConnection(d, drop)
def forward(self, x, mem, tgt_mask, src_mask):
x = self.s1(x, lambda y: self.self_a(y, tgt_mask)[0])
x = self.s2(x, lambda y: self.cross_a(y, src_mask, kv=mem)[0])
return self.s3(x, self.ff)
class DateTransformer(nn.Module):
def __init__(self, vocab, d=64, h=4, n_enc=2, n_dec=2, drop=0.1):
super().__init__()
self.src_embed = TokenAndPositional(vocab, d)
self.tgt_embed = TokenAndPositional(vocab, d)
self.enc = nn.ModuleList([EncoderLayer(d, h, drop) for _ in range(n_enc)])
self.dec = nn.ModuleList([DecoderLayer(d, h, drop) for _ in range(n_dec)])
self.norm = nn.LayerNorm(d)
self.head = nn.Linear(d, vocab)
def encode(self, src, src_mask=None):
x = self.src_embed(src)
for layer in self.enc:
x = layer(x, src_mask)
return self.norm(x)
def decode(self, tgt, mem, tgt_mask, src_mask=None):
y = self.tgt_embed(tgt)
for layer in self.dec:
y = layer(y, mem, tgt_mask, src_mask)
return self.norm(y)
def forward(self, src, tgt, tgt_mask, src_mask=None):
mem = self.encode(src, src_mask)
y = self.decode(tgt, mem, tgt_mask, src_mask)
return self.head(y)
The MultiHeadAttention class from module 3 needs one small extension to accept a separate kv input for cross-attention. That extension is a five-line change: split qkv into q from x and k, v from kv when the argument is given.
Shape unit tests before training
The single most common bug in a hand-built Transformer is a shape mismatch that trains anyway with silent broadcasting. Assert shapes explicitly before you launch anything long.
def test_shapes():
m = DateTransformer(len(VOCAB), d=64, h=4, n_enc=2, n_dec=2)
src = torch.zeros(3, 20, dtype=torch.long)
tgt = torch.zeros(3, 12, dtype=torch.long)
tgt_mask = torch.tril(torch.ones(12, 12)).view(1, 1, 12, 12)
out = m(src, tgt, tgt_mask)
assert out.shape == (3, 12, len(VOCAB)), out.shape
print("shapes OK")
test_shapes()
Total parameter count on the values above: about 130k. That is small enough to fit on a laptop CPU and train in ten minutes.
Training loop
from torch.optim import AdamW
import torch.nn.functional as F
def train(model, steps=3000, batch=64, lr=3e-4):
loader = DataLoader(DateDataset(20000), batch_size=batch, shuffle=True)
opt = AdamW(model.parameters(), lr=lr)
it = iter(loader)
for step in range(steps):
try:
src, tgt = next(it)
except StopIteration:
it = iter(loader); src, tgt = next(it)
tgt_in, tgt_out = tgt[:, :-1], tgt[:, 1:]
T = tgt_in.size(1)
tgt_mask = torch.tril(torch.ones(T, T)).view(1, 1, T, T)
logits = model(src, tgt_in, tgt_mask)
loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)),
tgt_out.reshape(-1))
opt.zero_grad(); loss.backward(); opt.step()
if step % 200 == 0:
print(f"step {step} loss {loss.item():.3f}")
Loss should drop from about 3.9 (random over 65-symbol vocab) to below 0.05 within two thousand steps. Anything else points to a bug — most often the missing causal mask, a mis-shifted target, or padding not being ignored in the loss.
Reading the attention maps
The whole reason we picked this task: attention weights become interpretable. After training, feed a single example and plot the cross-attention weights averaged over heads.
import matplotlib.pyplot as plt
def plot_attention(model, src_str, tgt_str):
model.eval()
src = torch.tensor(encode(src_str.ljust(20))).unsqueeze(0)
tgt = torch.tensor(encode(SOS + tgt_str)).unsqueeze(0)
mem = model.encode(src)
y = model.tgt_embed(tgt)
_, w = model.dec[-1].cross_a(y, kv=mem) # last layer, mean over heads
w = w.mean(dim=1).squeeze(0).detach()
plt.imshow(w, aspect="auto"); plt.colorbar()
plt.yticks(range(len(tgt_str) + 1), SOS + tgt_str)
plt.xticks(range(20), list(src_str.ljust(20)), rotation=90)
plt.title("cross-attention weights")
plt.show()
plot_attention(model, "3 March 2026", "2026-03-03")
You should see three bright regions: the four year digits in the output looking at the four year digits in the input, the month digits looking at "March", and the day digits looking at the leading "3". If those regions are misaligned or diffuse, either training has not converged or one of your building blocks has a bug. The map is the test.
Where this Transformer stops
The date task is small enough to train from scratch. Anything larger — sentiment, translation, code generation — needs pre-trained weights, which is what modules 6, 7 and 8 introduced. Our two hundred lines of PyTorch reach the same architecture as those models; the difference is that a real model saw a trillion tokens during pre-training, and ours saw twenty thousand.
If you want to push further, three natural extensions live inside the same file:
- Swap sinusoidal positional encoding for RoPE (module 4). The maps stay just as readable.
- Replace our hand-rolled attention with
F.scaled_dot_product_attention(module 9) to enable FlashAttention. - Add a beam search decoding routine (module 7) for tasks where greedy decoding matters.
When loss stays flat, always test in this order: shapes, then a single-example forward pass, then a fixed micro-batch overfitting run. A working Transformer overfits a batch of four examples to loss below 0.01 within five hundred steps. If yours does not, the bug is in a block you already wrote, not in the hyperparameters.
Commit the code at the end of each red-thread module in a separate file: v1_single_head.py, v2_multi_head.py, v3_with_pos_enc.py, and so on. When something breaks, you can bisect exactly which module's changes caused it.
In summary
- The full Transformer is the assembly of every block from modules 2 to 8, plus attention efficiency choices from module 9, and fits in about two hundred lines of PyTorch.
- The date translation task is small enough to train on a laptop and produces attention maps whose brightness lines up with the source positions — a visible proof that every block works.
- Shape assertions catch the majority of hand-built Transformer bugs before training; a working model overfits a tiny batch to near-zero loss in a few hundred steps.
- Extensions — RoPE, FlashAttention, beam search — plug into the same file without touching the core, which is why we spent so long on modular design.
Next module: the recap and the 40-question exam that ends the course.