Build the GPT

Notebook 08 · Assemble everything

Stack the blocks, add embeddings and a head -- that’s a GPT.

A guided tour of gpt.py: every class, what it does, and why. Every line traces back to something you built by hand in Part 1.

This notebook needs PyTorch and a GPU A browser can't run PyTorch, so the cells below are shown with the outputs of a real run (Apple-GPU). Download the notebook at the bottom of the page to run and train it yourself.

Now we assemble everything (attention, multi-head attention, the feed-forward network, LayerNorm, the residuals, and the embeddings) into a complete GPT. To keep it reusable, the model lives in a file called gpt.py, so that notebooks 09 and 10 can simply import it. This notebook is a guided tour of that file: we read each class in turn and explain what it does and why.

Here is the whole journey a token takes, from top to bottom:

python

token ids ->  token embedding  +  position embedding
          ->  N x Transformer block   (attention + feed-forward, with residuals)
          ->  final LayerNorm
          ->  linear head  ->  logits over the vocabulary (next-token scores)

In words: turn each token id into a vector and add in where it sits; let the stack of blocks mix and refine those vectors; give them a final clean-up; then read out a score for every possible next token. That last list of scores is exactly what softmax and sampling turn into the next character.

token ids(batch, time)

▼

token embedding + position embeddingids become vectors that carry meaning and order

▼

N × Transformer blockattention + feed-forward (notebook 06)

▼

final LayerNorm → linear head

▼

logits: a score for every next token(batch, time, vocab)

python

# Colab setup -- fetch the files this notebook needs.
# (Does nothing when run locally in the course folder.)
import os, urllib.request
BASE = ("https://raw.githubusercontent.com/waze"
        "emlabs/llm-book-code/main/")
for f in ['gpt.py', 'data/input.txt']:
    if not os.path.exists(f):
        d = os.path.dirname(f)
        if d: os.makedirs(d, exist_ok=True)
        urllib.request.urlretrieve(BASE + f, f)
        print("downloaded", f)

Line by line: what each line does

import os, urllib.request: two standard-library toolboxes: checking files on disk, and downloading from the web.
BASE = (...): the web address of this course's folder on GitHub, split over two lines (Python glues adjacent strings together).
for f in [...]: loop over the file names this notebook needs.
if not os.path.exists(f): only download what is missing -- running locally, everything already exists, so nothing happens.
os.makedirs(d, exist_ok=True): create the folder for the file if it has one (exist_ok means don't complain if it's already there).
urllib.request.urlretrieve(...): download the file and save it under the same name here.

python

import torch, inspect
from gpt import GPTConfig, Head, MultiHeadAttention, FeedForward, Block, GPT
torch.manual_seed(1337)                       # reproducible init + untrained sample
device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
print("model lives in gpt.py | device:", device)

Line by line: what each line does

import torch, inspect: PyTorch, plus inspect, a standard tool that can fetch a class's source code as text.
from gpt import GPTConfig, Head, ...: import the model's classes from the local file gpt.py, so notebooks 08-10 all share the exact same code.
torch.manual_seed(1337): fix the seed so the untrained sample lower down is reproducible (the same seed the other notebooks use).
The device line picks Apple GPU → NVIDIA → CPU, as before.

outputmodel lives in gpt.py | device: mps

Config

Every size of the model, gathered in one place. Think of these as the dials you turn up for a larger and more capable (but slower and more demanding) model: more layers (n_layer) makes it deeper; a wider embedding (n_embd) gives each token more room to represent meaning; more heads (n_head) lets it track more relationships at once; and a longer block_size lets it look further back. Our values are deliberately tiny, so the whole thing trains in about a minute.

python

print(inspect.getsource(GPTConfig))

Line by line: what each line does

print(inspect.getsource(GPTConfig)): fetch and print the class's real source from gpt.py. What you see below is the actual code.
@dataclass: a decorator (the @ line above the class) that auto-writes the boring constructor: you just declare fields with types and defaults, and GPTConfig(vocab_size=65) works.
vocab_size: int = 65: a field: its name, its type, its default. Pure settings, no behavior. The other fields (block_size, n_layer, n_head, n_embd) are the size dials you turn up for a bigger model.

output@dataclass class GPTConfig: vocab_size: int = 65 # number of distinct tokens block_size: int = 128 # maximum context length (tokens the model can look back over) n_layer: int = 4 # number of Transformer blocks n_head: int = 4 # attention heads per block n_embd: int = 128 # embedding / residual-stream width dropout: float = 0.0 # dropout probability (set >0 to regularize larger runs)

One attention head

This is identical to notebook 05, now written as a trainable nn.Module. key, query, and value are the learned projections; the tril buffer enforces causality (no looking at the future); and the 1/sqrt(d) scaling keeps softmax soft (notebook 05, step 4). Nothing here is new; it is the same head, wired up so that autograd can train it.

python

print(inspect.getsource(Head))

Line by line: what each line does

This prints notebook 07's Head, with two production touches:
self.key/query/value = nn.Linear(...): the three learned projections, as layers.
self.register_buffer("tril", ...): the causal mask, carried with the model but not trained.
self.dropout = nn.Dropout(cfg.dropout): during training, randomly zero a fraction of the attention weights (regularization, explained in notebook 09). With dropout=0.0 it does nothing.
wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5 ... -- scale, mask, softmax, blend: line-for-line what you built in notebooks 05 and 07.
The """...""" right under the class line is a docstring -- documentation tools can read it; Python ignores it at runtime.

outputclass Head(nn.Module): """One head of causal self-attention (notebook 05, now with learned, trainable weights).""" def __init__(self, cfg: GPTConfig, head_size: int): super().__init__() self.key = nn.Linear(cfg.n_embd, head_size, bias=False) self.query = nn.Linear(cfg.n_embd, head_size, bias=False) self.value = nn.Linear(cfg.n_embd, head_size, bias=False) # a constant lower-triangular mask; registered as a buffer so it moves with .to(device) self.register_buffer("tril", torch.tril(torch.ones(cfg.block_size, cfg.block_size))) self.dropout = nn.Dropout(cfg.dropout) def forward(self, x): B, T, C = x.shape k = self.key(x) # (B, T, head_size) q = self.query(x) # (B, T, head_size) # scaled dot-product affinities, then mask out the future wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5 # (B, T, T) wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf")) wei = F.softmax(wei, dim=-1) wei = self.dropout(wei) v = self.value(x) # (B, T, head_size) return wei @ v # (B, T, head_size)

Multi-head attention

Run several heads in parallel and join their outputs together, so the model can attend to different kinds of relationship at once, then use a projection (proj) to mix them back to the embedding width (notebook 06). The split head_size = n_embd // n_head is what keeps the heads affordable: four heads each do a quarter-width job, so the total cost matches that of a single full-width head.

python

print(inspect.getsource(MultiHeadAttention))

Line by line: what each line does

head_size = cfg.n_embd // cfg.n_head: split the width across heads (128 / 4 = 32 each).
self.heads = nn.ModuleList([Head(cfg, head_size) for _ in range(cfg.n_head)]): a list of heads that PyTorch knows about (a plain Python list would hide their parameters from the optimizer).
torch.cat([h(x) for h in self.heads], dim=-1): run every head, then glue their outputs side by side along the last axis: 4 heads × 32 = 128. This is the "concatenate" from notebook 06.
self.proj then self.dropout: a final projection that mixes the heads, then dropout.

outputclass MultiHeadAttention(nn.Module): """Several attention heads in parallel, concatenated and projected (notebook 06).""" def __init__(self, cfg: GPTConfig): super().__init__() head_size = cfg.n_embd // cfg.n_head self.heads = nn.ModuleList([Head(cfg, head_size) for _ in range(cfg.n_head)]) self.proj = nn.Linear(cfg.n_embd, cfg.n_embd) self.dropout = nn.Dropout(cfg.dropout) def forward(self, x): out = torch.cat([h(x) for h in self.heads], dim=-1) # (B, T, n_embd) return self.dropout(self.proj(out))

Feed-forward

After the tokens have traded information through attention, each one is processed on its own by a small MLP: expand to four times the width, apply GELU, then project back (notebook 06). This is the "now think about what you just heard" step, and most of the model's numbers live here.

python

print(inspect.getsource(FeedForward))

Line by line: what each line does

self.net = nn.Sequential(...): a pipeline: each layer feeds the next, in order. Calling self.net(x) runs the whole chain.
nn.Linear(n_embd, 4 * n_embd) → nn.GELU() → nn.Linear(4 * n_embd, n_embd): widen 4×, bend, narrow back: notebook 06's feed-forward, with PyTorch's exact GELU.
nn.Dropout(cfg.dropout): the same regularization knob at the end.
def forward(self, x): return self.net(x): just run the pipeline.

outputclass FeedForward(nn.Module): """Per-token MLP: expand 4x, nonlinearity, project back. Where most parameters live.""" def __init__(self, cfg: GPTConfig): super().__init__() self.net = nn.Sequential( nn.Linear(cfg.n_embd, 4 * cfg.n_embd), nn.GELU(), nn.Linear(4 * cfg.n_embd, cfg.n_embd), nn.Dropout(cfg.dropout), ) def forward(self, x): return self.net(x)

The block

One block is attention (communicate) followed by the feed-forward network (think), each wrapped in pre-norm and a residual connection. Stacking these blocks is the entire depth of the network. The residual, written x = x + sublayer(norm(x)) (keep the input, add an edit on top), is what lets the gradients flow cleanly through a deep stack (notebook 06).

python

print(inspect.getsource(Block))

Line by line: what each line does

self.ln1 = nn.LayerNorm(cfg.n_embd) (and ln2) -- the LayerNorm you wrote by hand in notebook 06, now built in, with learned dials.
self.sa = MultiHeadAttention(cfg) / self.ffwd = FeedForward(cfg): the two sub-layers.
x = x + self.sa(self.ln1(x)): normalize → attend → add back (residual).
x = x + self.ffwd(self.ln2(x)): normalize → think → add back. Two lines, the whole block.

outputclass Block(nn.Module): """Transformer block: communicate (attention), then compute (MLP). Pre-norm + residuals.""" def __init__(self, cfg: GPTConfig): super().__init__() self.ln1 = nn.LayerNorm(cfg.n_embd) self.sa = MultiHeadAttention(cfg) self.ln2 = nn.LayerNorm(cfg.n_embd) self.ffwd = FeedForward(cfg) def forward(self, x): x = x + self.sa(self.ln1(x)) # residual connection around attention x = x + self.ffwd(self.ln2(x)) # residual connection around the feed-forward return x

The full GPT

This is the whole model, and every line should look familiar by now:

token_embedding: a lookup table that turns each token id into a learnable vector (its "meaning"). Nothing assigns these meanings by hand; the vectors start out random and are nudged by training until tokens used in similar ways end up with similar vectors. Meaning emerges from usage.
position_embedding: a second table that adds in where the token sits (notebook 06's position information, learned here rather than fixed as sinusoids).
blocks: the stack of Transformer blocks that do the real work.
ln_f and then lm_head: a final LayerNorm clean-up, followed by a linear layer that turns each position's vector into logits over the vocabulary, that is, a score for every possible next token.

If you pass in targets, the model also returns the cross-entropy loss (notebook 01). And generate is the autoregressive sampling loop from notebook 03, with two new dials we will explore in notebook 10 (temperature and top_k): predict, take the last position's logits, optionally filter them, apply softmax, draw one token, append it, and repeat, always cropping the context to the last block_size tokens so the model never looks back further than it was trained to.

python

print(inspect.getsource(GPT))

Line by line: what each line does

self.token_embedding = nn.Embedding(cfg.vocab_size, cfg.n_embd): the meaning table: 65 rows (one per character), 128 numbers each.
self.position_embedding = nn.Embedding(cfg.block_size, cfg.n_embd): the same idea for positions 0..127.
self.blocks = nn.Sequential(*[Block(cfg) for _ in range(cfg.n_layer)]): build n_layer Blocks; the * unpacks the list into separate arguments.
self.ln_f / self.lm_head: a final LayerNorm, then a Linear that turns each token's 128 numbers into 65 next-token scores.
self.apply(self._init_weights): run a small setup on every sub-module: start weights small-random (std 0.02), biases at zero. Good starting points matter.
x = tok + pos: meaning plus position, added together (broadcasting fills the batch axis).
@torch.no_grad() above generate: a decorator meaning "don't record gradients": generation is pure inference.
Inside generate: crop to the last block_size tokens, keep only the last position's logits, divide by temperature, optionally keep just the top-k (torch.topk finds them; everything below gets -inf), softmax, then torch.multinomial(probs, 1) rolls the weighted die and torch.cat appends it.
num_params(): add up p.numel() (element count) over every weight grid: the parameter total.

outputclass GPT(nn.Module): """The full model: embeddings -> stacked blocks -> final norm -> next-token logits.""" def __init__(self, cfg: GPTConfig): super().__init__() self.cfg = cfg self.token_embedding = nn.Embedding(cfg.vocab_size, cfg.n_embd) self.position_embedding = nn.Embedding(cfg.block_size, cfg.n_embd) self.blocks = nn.Sequential(*[Block(cfg) for _ in range(cfg.n_layer)]) self.ln_f = nn.LayerNorm(cfg.n_embd) self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None): B, T = idx.shape tok = self.token_embedding(idx) # (B, T, n_embd) pos = self.position_embedding(torch.arange(T, device=idx.device)) # (T, n_embd) x = tok + pos # add position info x = self.blocks(x) # (B, T, n_embd) x = self.ln_f(x) logits = self.lm_head(x) # (B, T, vocab_size) loss = None if targets is not None: B, T, V = logits.shape loss = F.cross_entropy(logits.view(B * T, V), targets.view(B * T)) return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): """Autoregressively sample max_new_tokens, one token at a time.""" for _ in range(max_new_tokens): idx_cond = idx[:, -self.cfg.block_size:] # never look back further than block_size logits, _ = self(idx_cond) logits = logits[:, -1, :] / temperature # focus on the last position if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float("inf") probs = F.softmax(logits, dim=-1) idx_next = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, idx_next), dim=1) return idx def num_params(self): return sum(p.numel() for p in self.parameters())

Instantiate it

Build a real, if small, GPT sized for our character data, and count its numbers. It comes out at around 825,000, which sounds like a lot until you learn that GPT-3 has 175 billion, roughly 200,000 times more. The striking part is that the design you are looking at is, line for line, what the giant models use. They are the same thing, only far wider, far deeper, and trained on far more text.

824,897

parameters

layers

attention heads

128

context length

python

from pathlib import Path
text = Path("data/input.txt").read_text()
vocab_size = len(sorted(set(text)))

cfg = GPTConfig(vocab_size=vocab_size, block_size=128, n_layer=4, n_head=4, n_embd=128)
model = GPT(cfg).to(device)
print(f"vocab={cfg.vocab_size}  block={cfg.block_size}  layers={cfg.n_layer}  heads={cfg.n_head}  width={cfg.n_embd}")
print(f"total parameters: {model.num_params():,}")

Line by line: what each line does

cfg = GPTConfig(vocab_size=vocab_size, block_size=128, n_layer=4, n_head=4, n_embd=128): fill in the size dials.
model = GPT(cfg).to(device): build the model and ship every weight to the GPU.
f"{model.num_params():,}": the :, format inserts thousands separators: 824,897.

outputvocab=65 block=128 layers=4 heads=4 width=128 total parameters: 824,897

Where the 824,897 parameters live

That single number is easier to trust once you see it split up. Group every parameter tensor by the part of the model it belongs to, and one thing stands out: the stacked Transformer blocks hold about 96% of the weights, while the embeddings and the output head are a thin shell around them.

python

from collections import OrderedDict
groups = OrderedDict()
for name, p in model.named_parameters():
    top = name.split(".")[0]
    label = {"token_embedding": "token embedding", "position_embedding": "position embedding",
             "blocks": "4 transformer blocks", "ln_f": "final LayerNorm",
             "lm_head": "output head"}.get(top, top)
    groups[label] = groups.get(label, 0) + p.numel()

total = sum(groups.values())
for label, n in groups.items():
    print(f"  {label:22s} {n:8,d}  ({100*n/total:4.1f}%)")
print(f"  {'TOTAL':22s} {total:8,d}")

import matplotlib.pyplot as plt
plt.figure(figsize=(6.5, 3))
bars = plt.bar([l.replace(" ", "\n", 1) for l in groups], list(groups.values()),
               color=["#5b5bd6" if "blocks" in l else "#8b8bd6" for l in groups])
for b, n in zip(bars, groups.values()):
    plt.text(b.get_x() + b.get_width() / 2, n, f"{n:,}", ha="center", va="bottom", fontsize=7)
plt.ylabel("parameters"); plt.title("where the parameters live (blocks dominate)")
plt.xticks(fontsize=7); plt.tight_layout(); plt.show()

Line by line: what each line does

The loop reads each parameter tensor's name (like blocks.0.sa...) and adds its size to the matching component.
The token and position embeddings and the output head are small -- a few thousand numbers each; the four blocks together hold 791,552, about 96%.
Inside each block most of that is the feed-forward network's two big matrices (the 4× expand and contract from notebook 06).
The bar chart makes the imbalance plain: real capacity lives in the stacked blocks, which is why "make it bigger" almost always means more or wider blocks.

output token embedding 8,320 ( 1.0%) position embedding 16,384 ( 2.0%) 4 transformer blocks 791,552 (96.0%) final LayerNorm 256 ( 0.0%) output head 8,385 ( 1.0%) TOTAL 824,897 output figure

Forward pass and an untrained sample

A forward pass returns logits with shape (batch, time, vocab) and, when given targets, a loss near ln(vocab) = 4.17. This is because a freshly built model has random weights and is essentially guessing all 65 characters as equally likely (notebook 01's random baseline). Generating from it produces pure gibberish. That is our starting line; notebook 09 trains it into something that writes Shakespeare.

python

import math
B, T = 4, cfg.block_size
x = torch.randint(0, vocab_size, (B, T), device=device)
y = torch.randint(0, vocab_size, (B, T), device=device)
logits, loss = model(x, y)
print("logits:", tuple(logits.shape), "| loss: %.3f | ln(vocab) = %.3f" % (loss.item(), math.log(vocab_size)))

stoi = {c:i for i,c in enumerate(sorted(set(text)))}; itos = {i:c for c,i in stoi.items()}
decode = lambda l: "".join(itos[i] for i in l)
ctx = torch.zeros((1,1), dtype=torch.long, device=device)
print("\n--- untrained sample (gibberish) ---")
print(decode(model.generate(ctx, 200)[0].tolist()))

Line by line: what each line does

x = torch.randint(0, vocab_size, (B, T), device=device): random token ids shaped like a real batch, just to exercise the forward pass.
logits, loss = model(x, y): run the model; with targets it also returns the loss.
math.log(vocab_size): ln(65) = 4.174, the know-nothing baseline; the untrained loss lands right next to it.
ctx = torch.zeros((1,1), dtype=torch.long, device=device): the generation prompt: one sequence holding one token (id 0, the newline) -- a neutral start.
decode(model.generate(ctx, 200)[0].tolist()): generate 200 tokens, take the first (only) sequence, turn the tensor into a plain list, decode to text. Untrained = gibberish.

outputlogits: (4, 128, 65) | loss: 4.197 | ln(vocab) = 4.174 --- untrained sample (gibberish) --- XwO AOdTZYN3cESOPkTIlTqNGXy:.CeiZoFkci-:jP?vV?S?.MpsHhkYKrAMr hFulHerHEe.-o-XwKbQRT!dj!SCB$l n'm?'c3hhCZTV?koIuovIR OTCGD;PFh!VSBBj3UpFUapTSDDaZsQFFQiXNQT&:Nlp-$iMGbe SS McDkoSkjpvvpeQSxmn-FnFQajfL:S,

Follow the shapes through one forward pass

The diagram at the top of this page traced the path from token ids to logits. Here it is with real tensors, so you can watch each stage keep the batch and time axes and change only the last one: ids become 128-wide vectors, the blocks refine them without reshaping, and the head finally opens each position out to one score per vocabulary character.

Interactive · follow the shapes

One tensor, all the way through

The one thing that trips people up is shapes. Here is a real tensor flowing through the model. Notice the batch and time axes stay put the whole way — the model just reshapes each position's vector, ending with one score per vocab character.

python

# Run the model's stages by hand and print the shape after each (mirrors GPT.forward).
idx = torch.randint(0, cfg.vocab_size, (1, 20), device=device)
tok = model.token_embedding(idx)
pos = model.position_embedding(torch.arange(20, device=device))
h = tok + pos
hb = model.blocks(h)
hn = model.ln_f(hb)
logits = model.lm_head(hn)
print(f"idx (token ids)        : {tuple(idx.shape)}           (batch, time)")
print(f"token + position embed : {tuple(h.shape)}    (batch, time, width 128)")
print(f"after the 4 blocks     : {tuple(hb.shape)}    (same shape -- blocks preserve it)")
print(f"after final LayerNorm  : {tuple(hn.shape)}    (same shape)")
print(f"logits (lm_head)       : {tuple(logits.shape)}     (a score per vocab char, per position)")

Line by line: what each line does

idx: a batch of token ids, shape (batch, time).
token_embedding + position_embedding: each id becomes a 128-wide vector and the position vector is added -- now (batch, time, 128).
model.blocks(h): the four Transformer blocks run and keep the shape exactly, just refining the numbers.
ln_f then lm_head: a final normalize, then a linear layer opens each position out to one score per vocabulary character -- (batch, time, 65).
Every stage keeps the batch and time axes; only the last axis changes. That regularity is exactly what lets the blocks stack.

outputidx (token ids) : (1, 20) (batch, time) token + position embed : (1, 20, 128) (batch, time, width 128) after the 4 blocks : (1, 20, 128) (same shape -- blocks preserve it) after final LayerNorm : (1, 20, 128) (same shape) logits (lm_head) : (1, 20, 65) (a score per vocab char, per position)

Recap

The GPT is assembled and lives in gpt.py: embeddings, then blocks, then a final norm, then logits, with about a million numbers in total, and both the forward pass and generate already working (just untrained so far). Every line traces back to something you built by hand in Part 1; nothing here is new, it has only been wired together.

Next, notebook 09 trains it on Shakespeare and we watch the loss fall and the text sharpen.

⬇ Download this lesson as a notebook — 08_build_gpt.ipynb