GPT-2 Small · decoder stack (12× blocks, Pre-LN, learned wpe)
This chapter evolves the chapter 01 Transformer into GPT-2 — a decoder-only stack with learned positional embeddings, pre-norm residuals, GELU feed-forward blocks, and causal self-attention. Hyperparameters follow GPT-2 Small (117M) from Radford et al., matching the diagram above. We keep the same teaching pattern: one config object up front, then each layer in order until the full model runs.
Chapter 01 is encoder–decoder for seq2seq; this chapter keeps only the decoder path for next-token language modeling at GPT-2 Small scale (~117M parameters).
Each section below is one building block. You get a short explanation, the code for that module, a delta from chapter 01 where it helps, and how it connects to what came before. By the end you have a runnable GPT class, not just a diagram.
Delta from Transformer (01)
- Encoder removed — decoder-only (no cross-attention)
- Sinusoidal PE → learned positional embedding (
nn.Embedding) - Post-norm → pre-norm (LayerNorm before sub-blocks)
- ReLU → GELU activation in the FFN
- Output projection uses a linear (no
log_softmaxat the model boundary)
Full side-by-side with chapter 01: §08 Comparison.
One config object for the whole model
Before any module, we define a single GPTConfig dataclass for GPT-2 Small (117M). Every class below takes config instead of a long list of constructor arguments — same pattern as chapter 01, with fields adjusted for a decoder-only language model.
The GPT-2 paper defines four sizes (Small → XL). We use the Small variant to match the architecture figure on this page: d_model=768, N=12 blocks, H=12 heads, vocab=50,257, context length 1,024, dropout=0.1, and d_ff=4× d_model (~117M parameters with weight tying).
dk property: dmodel // H → 768 // 12 = 64. Multi-head attention splits d_model across H heads; d_model must be divisible by H.
dff property: 4 * dmodel → 3072 — the FFN inner dimension used throughout the stack.
All values below are GPT-2 Small (117M) — change GPTConfig fields to explore other sizes from the paper.
dmodel768 — n_embd / hidden_size (GPT-2 Small)vocabsize50,257 — BPE vocabulary sizeB2 — batch size for demosmax_seq_len1024 — context lengthN12 — n_layer (decoder blocks)H12 — n_headdropout0.1 — dropout probabilitydk64 — dmodel // Hdff3072 — 4 × dmodel (n_inner)1from dataclasses import dataclass
2import torch
3import torch.nn as nn
4import torch.nn.functional as F
5import math
6
7@dataclass
8class GPTConfig:
9 dmodel: int = 768 # GPT-2 Small (117M)
10 vocabsize: int = 50257 # BPE vocabulary size
11 B: int = 2 # batch size (for demo runs)
12 max_seq_len: int = 1024 # context length
13 N: int = 12 # number of transformer blocks
14 H: int = 12 # number of attention heads
15 dropout: float = 0.1 # dropout probability
16
17 @property
18 def dk(self):
19 # per-head dimension — dmodel must be divisible by H
20 return self.dmodel // self.H
21
22 @property
23 def dff(self):
24 return 4 * self.dmodel
25
26
27config = GPTConfig() # GPT-2 Small (117M) — shared by every module belowPosition as a learnable table
Instead of fixed sinusoids, GPT-2 learns a matrix of shape (max_seq_len, dmodel). wte maps token IDs to vectors; wpe maps position indices (0..S-1) to vectors. Both are nn.Embedding. Their outputs are summed then dropped — that combined embedding is the only input the stack sees.
Chapter 01 wrapped these steps in separate classes (InputEmbedding, PositionalEncoding). Here we skip that — you already know how to build a module class, and these two tables are just nn.Embedding lookups. So we declare them directly on GPT as self.wte and self.wpe, then sum them in forward. Same math, less ceremony.
1# inside GPT.__init__ — two learned embedding tables:
2self.wte = nn.Embedding(config.vocabsize, config.dmodel) # wte(x): (B, S, dmodel)
3self.wpe = nn.Embedding(config.max_seq_len, config.dmodel) # wpe(pos): (S, dmodel)
4
5# inside GPT.forward — sum and drop:
6_, S = x.shape # x: (B, S)
7pos = torch.arange(0, S, device=x.device) # (S,)
8x = self.drop(self.wte(x) + self.wpe(pos)) # (B, S, dmodel) + (S, dmodel) → (B, S, dmodel)The 2017 Transformer’s PositionalEncoding: a fixed pe table built from sines and cosines, added to the scaled token embedding (excerpt from part 01).
1 def forward(self, x):
2 x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
3 return self.dropout(x)Same LayerNorm, different place in the block
We already built LayerNormalization from scratch in chapter 01 — here we use PyTorch’s built-in nn.LayerNorm directly. Same normalize–scale–shift math; the practical shift is where you apply it: before each sublayer (pre-norm), not only after a residual (post-norm).
No custom class needed — nn.LayerNorm(d_model) normalizes the last dimension for each position.
1# Drop-in replacement for the educational LayerNormalization class: same math,
2# fused kernels, and per-feature learnable weight/bias in modern PyTorch.
3ln = nn.LayerNorm(config.dmodel) # (B, S, dmodel) → (B, S, dmodel)Expand, GELU, contract
Same two-linear pattern as the 2017 FFN: dmodel → dff → dmodel (dff = 4 × dmodel). The nonlinearity is GELU (Gaussian error linear unit), a smooth mix of identity and ReLU. Dropout is applied after the activation, between the two linears. Named MLP in our code — no standalone FFN class needed.
1class MLP(nn.Module):
2 def __init__(self, config):
3 super().__init__()
4 self.w1 = nn.Linear(config.dmodel, config.dff) # (B, S, dmodel) → (B, S, dff)
5 self.w2 = nn.Linear(config.dff, config.dmodel) # (B, S, dff) → (B, S, dmodel)
6 self.dropout = nn.Dropout(config.dropout)
7
8 def forward(self, x):
9 # x: (B, S, dmodel) → w1 → gelu → dropout → w2 → (B, S, dmodel)
10 return self.w2(self.dropout(F.gelu(self.w1(x)))) # (B, S, dmodel)Feed-forward in part 01 used ReLU (and often dropout) between the two linears.
1 def forward(self, x):
2 return self.linear_2(self.dropout(torch.relu(self.linear_1(x))))Causal self-attention — Wq, Wk, Wv
Same projection layout as chapter 01: separate Wq, Wk, Wv, and Wo matrices. For causal self-attention, Q, K, and V all come from the same sequence x. GPT-2 adds a causal mask (sliced to [:S, :S] at runtime) and two dropouts: attn_dropout on attention weights, resid_dropout on the output.
1class CausalSelfAttention(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 self.dmodel = config.dmodel
5 self.H = config.H
6 self.dk = config.dk
7 self.wq = nn.Linear(config.dmodel, config.dmodel) # dmodel → dmodel
8 self.wk = nn.Linear(config.dmodel, config.dmodel) # dmodel → dmodel
9 self.wv = nn.Linear(config.dmodel, config.dmodel) # dmodel → dmodel
10 self.wo = nn.Linear(config.dmodel, config.dmodel) # dmodel → dmodel
11 self.attn_dropout = nn.Dropout(config.dropout) # on attention weights
12 self.resid_dropout = nn.Dropout(config.dropout) # on output residual
13
14 def attention(self, q, k, v, S, mask):
15 # q, k, v: (B, H, S, dk)
16 attn_score = q @ k.transpose(-1, -2) / math.sqrt(self.dk) # (B, H, S, S)
17 attn_score = attn_score.masked_fill(mask[:S, :S] == 0, -1e9) # mask: (S, S)
18 attn_score = attn_score.softmax(dim=-1) # (B, H, S, S)
19 attn_score = self.attn_dropout(attn_score) # (B, H, S, S)
20 return attn_score @ v # (B, H, S, dk)
21
22 def forward(self, x, mask):
23 B, S, _ = x.shape # x: (B, S, dmodel)
24 # project and split into H heads
25 query = self.wq(x).view(B, S, self.H, self.dk).transpose(1, 2) # (B, H, S, dk)
26 key = self.wk(x).view(B, S, self.H, self.dk).transpose(1, 2) # (B, H, S, dk)
27 value = self.wv(x).view(B, S, self.H, self.dk).transpose(1, 2) # (B, H, S, dk)
28 x = self.attention(query, key, value, S, mask) # (B, H, S, dk)
29 x = x.transpose(1, 2).contiguous().view(B, S, self.dmodel) # (B, S, dmodel)
30 return self.resid_dropout(self.wo(x)) # (B, S, dmodel)Norm before each sublayer — inline in Block
GPT-2 uses pre-norm: apply LayerNorm to the input before each sublayer, then add the residual back. Our code does this inline inside Block.forward with two separate norms (norm_1 for attention, norm_2 for FFN) — no separate wrapper class needed. The dropout inside each sublayer handles regularization.
1# No wrapper class needed — Block does pre-norm directly in forward
2def forward(self, x, mask):
3 # x: (B, S, dmodel)
4 x = x + self.attn(self.norm_1(x), mask) # attn out: (B, S, dmodel)
5 x = x + self.ff(self.norm_2(x)) # ff out: (B, S, dmodel)
6 return x # (B, S, dmodel)The 2017 stack is often described as post-norm at the block output: add the residual, then normalize.
1 def forward(self, x, sublayer):
2 return self.norm(x + self.dropout(sublayer(x)))norm + attn + norm + FFN, no cross-attention
A GPT-2 block is Block(config): two LayerNorms (norm_1, norm_2), one CausalSelfAttention, one MLP. The norms and dropouts live inside the block itself — no wrapper class needed. There is no cross-attention; only self-attention over the running prefix.
1class Block(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 self.norm_1 = nn.LayerNorm(config.dmodel) # (B, S, dmodel) → (B, S, dmodel)
5 self.attn = CausalSelfAttention(config) # (B, S, dmodel) → (B, S, dmodel)
6 self.norm_2 = nn.LayerNorm(config.dmodel) # (B, S, dmodel) → (B, S, dmodel)
7 self.ff = MLP(config) # (B, S, dmodel) → (B, S, dmodel)
8
9 def forward(self, x, mask):
10 # x: (B, S, dmodel)
11 x = x + self.attn(self.norm_1(x), mask) # (B, S, dmodel)
12 x = x + self.ff(self.norm_2(x)) # (B, S, dmodel)
13 return x # (B, S, dmodel)The original DecoderBlock in part 01: masked self-attention, then cross-attention to the encoder, then FFN (three residual sublayers).
1 def forward(self, x, encoder_output, src_mask, tgt_mask):
2 x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))
3 x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output, encoder_output, src_mask))
4 x = self.residual_connections[2](x, self.feed_forward_block)
5 return xToken + position, N Blocks, final norm, lm_head
GPT owns everything. wte + wpe produce the summed embedding; drop regularises it. N Blocks via nn.ModuleList, each taking (x, mask). A final LayerNorm then lm_head projects to vocab. Two key design choices: (1) weight tying — lm_head.weight = wte.weight halves parameters; (2) a pre-registered causal mask that is sliced per sequence at runtime — no mask reallocation on every forward pass.
With the §00 defaults, that is the full GPT-2 Small (117M) stack at 768 / 12 / 12 / 3072.
1class GPT(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 self.wte = nn.Embedding(config.vocabsize, config.dmodel) # wte(x): (B, S, dmodel)
5 self.wpe = nn.Embedding(config.max_seq_len, config.dmodel) # wpe(pos): (S, dmodel)
6 self.drop = nn.Dropout(config.dropout)
7 self.blocks = nn.ModuleList([Block(config) for _ in range(config.N)])
8 self.norm = nn.LayerNorm(config.dmodel) # (B, S, dmodel) → (B, S, dmodel)
9 self.lm_head = nn.Linear(config.dmodel, config.vocabsize, bias=False) # (B, S, dmodel) → (B, S, vocabsize)
10 self.lm_head.weight = self.wte.weight # weight tying — lm_head shares weights with wte
11 # full causal mask — every token attends to all previous tokens
12 self.register_buffer("mask", torch.tril(torch.ones(config.max_seq_len, config.max_seq_len))) # (max_seq_len, max_seq_len)
13
14 def forward(self, x):
15 _, S = x.shape # x: (B, S)
16 pos = torch.arange(0, S, device=x.device) # (S,)
17 x = self.drop(self.wte(x) + self.wpe(pos)) # (B, S, dmodel) + (S, dmodel) → (B, S, dmodel)
18 for block in self.blocks: x = block(x, self.mask) # (B, S, dmodel)
19 x = self.norm(x) # (B, S, dmodel)
20 return self.lm_head(x) # (B, S, vocabsize)Transformer (01) vs GPT-2 Small (02)
Canonical comparison for this chapter (defaults: GPT-2 Small, 117M) — same pair as the delta card above, with more detail after you have read the stack.
| Axis | Original Transformer | GPT-2 Small |
|---|---|---|
| Architecture | Encoder + decoder, seq2seq | Decoder-only stack (12 layers, no encoder) |
| Attention | Encoder: full bidirectional · Decoder: causal self + cross to encoder | Masked self-attention only (no cross-attention) |
| Position encoding | Sinusoidal, added at input | Learned wpe, summed with token embeddings |
| FFN activation | ReLU (in the 01 build) | GELU |
| Normalization | Post-norm (after residual add) | Pre-norm (before each sub-block) |
| Parameters | ~65M (6+6 layers, d=512, 8 heads) | ~117M (12 decoder layers, d=768, H=12) |
| Training goal | Translation / seq2seq | Next-token prediction |
The base Transformer paper uses 6 encoder + 6 decoder layers; GPT-2 Small uses 12 decoder-only layers — same layer count, but every block is devoted to generation.
Papers & technical sources
Primary reports and references for this chapter. Read these for full equations, training details, and official hyperparameters.
- Language Models are Unsupervised Multitask Learners (Radford et al., 2019)GPT-2 technical report — decoder-only, learned positional embeddings, pre-norm.
- Hugging Face: openai-community/gpt2Official config.json field names for GPT-2 Small.
Comments