The Illustrated Decoder

in nanoGPT

tok_emb = self.transformer.wte(idx)      # ids -> vectors
x = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:         # the repeated block, ×12
    x = block(x)
x = self.transformer.ln_f(x)
logits = self.lm_head(x)                 # vectors -> vocab scores

This whole page is one short method. The for-loop is the transformer — everything else is get-in (embeddings) and get-out (final norm + head). Every block takes and returns the same shape (B,T,C), which is exactly why they stack.

model.py:170–193 · GPT.forward · gpt notes

in nanoGPT

for _ in range(max_new_tokens):
    idx_cond = idx[:, -self.config.block_size:]  # crop context to the window
    logits, _ = self(idx_cond)                   # one full forward pass
    ...                                          # pick one token — figure 9
    idx = torch.cat((idx, idx_next), dim=1)      # append, run again

The writing loop really is this small: forward, pick, append, repeat. The crop on the first line is the model's hard memory limit — only block_size (1024) position vectors exist, so once the text outgrows the window, the oldest tokens simply fall off.

model.py:306–330 · GPT.generate · gpt notes §6

in nanoGPT

tok_emb = self.transformer.wte(idx)  # token lookup table    (50304 × 768)
pos_emb = self.transformer.wpe(pos)  # position lookup table  (1024 × 768)
x = self.transformer.drop(tok_emb + pos_emb)

Both are lookups, not matmuls — nn.Embedding just selects a row, and training tunes what each row means. One honest difference from the figure: the sin/cos ladder is the classic 2017 recipe, while nanoGPT (like GPT-2) learns its position vectors — wpe starts random and converges by backprop. Same job either way: inject order into an otherwise order-blind stack.

model.py:177–179 · wte + wpe · gpt notes §2

in nanoGPT

# .view() / .transpose()  — relabel & re-file, no math
# @                       — the one op that actually computes
att = q @ k.transpose(-2, -1)   # (B,nh,T,hs) @ (B,nh,hs,T) -> (B,nh,T,T)

The way experts read PyTorch: track the shape after every line, not the words. In the whole attention class, @ appears exactly twice — the scores and the blend; every reshape around them moves no numbers. And shape-legal ≠ correct: the notes walk through how matmul-ing the wrong pair of axes runs without error and silently computes garbage.

model.py:67 · attention notes §0

in nanoGPT

self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)   # 768 -> 2304, sliced back

All three lenses are one Linear (768 → 2304) — one big matmul beats three small ones on a GPU — then .split slices q, k, v apart, pure bookkeeping. And a division of labor the figure can't show: q and k only ever decide how much; v alone decides what — once the weights are computed, q and k never appear again.

model.py:35 + model.py:56 · attention notes

in nanoGPT

att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
y = att @ v

Score, mask, softmax, blend — one line each. exp(−inf) = 0, so the softmax is where causality actually lands: future tokens get exactly zero weight. On PyTorch ≥ 2.0 nanoGPT replaces all four lines with one flash-attention call (line 64) — same math, computed with running totals so the T×T grid is never materialized.

model.py:61–71 · CausalSelfAttention · attention notes §3–4

in nanoGPT

k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
# ...attention runs on every head at once, as a batch dimension...
y = y.transpose(1, 2).contiguous().view(B, T, C)  # stitch heads back together

Heads are a reshape, not separate modules: the 768-wide tensor is relabeled into 12 slices of 64, the batched matmul treats each slice as its own little attention, and the outputs are stitched straight back. Across 12 layers × 12 heads, GPT-2 small runs 144 distinct readers, each with its own learned weights.

model.py:56–72 · attention notes §4

in nanoGPT

def forward(self, x):
    x = x + self.attn(self.ln_1(x))   # gather across tokens, ADD it back
    x = x + self.mlp(self.ln_2(x))    # think within each token, ADD it back
    return x

The entire transformer block is these two lines, stacked 12 times. Look where the norms sit (pre-LN): each wraps a sublayer's input, and the raw un-normalized x is what's added — nothing ever rescales the stream itself, and that clean highway is why 12–96-layer stacks train at all. The mlp being called is just widen → gelu → shrink — and those two fat matrices hold roughly two-thirds of the model's parameters.

model.py:94–106 · Block + model.py:78–92 · MLP · block notes · mlp notes

in nanoGPT

logits = logits[:, -1, :] / temperature              # last position, scaled
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)   # SAMPLE — not argmax

multinomial draws from the distribution instead of taking the max — that's why one prompt writes different text each run. And a quiet trick in the projection before it: nanoGPT ties wte and lm_head to one shared matrix (line 138) — the table that turned ids into vectors also scores vectors against every id. Same similarity question, asked in both directions.

model.py:306–330 · GPT.generate · gpt notes §4–6