Chapter 6 · Section 6.2 · 55 min

The model

In §5.6 you assembled a Transformer from blocks and counted its parameters as about $12Ld^2$ , where $L$ is the number of layers and $d$ the width. Plug in GPT-2 small, with $L = 12$ and $d = 768$ , and you get 85 million. But the checkpoint OpenAI released has 124,439,808 parameters. Where do the other 39 million come from, and why would a model spend almost a third of its budget on them?

This section answers that by building GPT-2 exactly, the real architecture with every detail that matters, first in NumPy so you can see each operation, then in PyTorch. The two will agree to within $10^{-15}$ , and the PyTorch version will count to 124,439,808 on the nose. Along the way we'll meet the three details that make GPT-2 more than "a stack of Transformer blocks", namely where the LayerNorms go, why the input and output embeddings are the same matrix, and why some weights start $\sqrt{2L}$ times smaller than the rest.

The whole model in one picture

A GPT maps a batch of token ids to a batch of next-token scores. Read the diagram top to bottom, with shapes for a batch of $B$ sequences of $T$ tokens, vocabulary size $V$ , width $d$ , and $h$ attention heads of size $d_k = d / h$ .

idx                         (B, T)        integers in [0, V)
  │  token embedding  E[idx]  +  position embedding  P[0:T]
x⁰                          (B, T, d)
  │  L times:   x ← x + Attn(LN₁(x))      causal self-attention, h heads
  │             x ← x + MLP(LN₂(x))       d → 4d → d with GELU
xᴸ                          (B, T, d)
  │  final LayerNorm, then multiply by Eᵀ (the same E as at the bottom)
logits                      (B, T, V)     softmax(logits[b, t]) predicts token t + 1

Here $\mathbf{E} \in \mathbb{R}^{V \times d}$ is the token embedding table, $\mathbf{P} \in \mathbb{R}^{T \times d}$ is the learned position table, and $T$ here is the maximum context, 1,024 for GPT-2. Every arrow is something you've built before. The new things are the exact choices.

	GPT-2 small	our tiny model
vocabulary $V$	50,257 (byte-level BPE)	56 (characters plus EOT)
context $T$	1,024	64
layers $L$	12	2
width $d$	768	64
heads $h$ , head size $d_k$	12, 64	4, 16
MLP width	3,072 ( $4d$ )	256 ( $4d$ )
parameters	124,439,808	107,776

Embeddings, tokens plus positions

The first layer looks up a row of $\mathbf{E}$ for each token, exactly the lookup table of §1.2, and adds a row of $\mathbf{P}$ for each position,

\mathbf{x}^0_t = \mathbf{E}_{w_t,:} + \mathbf{P}_{t,:}, \qquad t = 0, \dots, T-1 .

In plain words, the vector at position $t$ starts as "which token is here" plus "where here is". Attention without a mask can't tell positions apart (it treats its input as a set, §4.6). The causal mask leaks a little order, since position $t$ averages over exactly $t + 1$ inputs, and causal models trained with no position embedding at all do learn some positional information (Haviv et al., 2022). But that signal is weak and indirect, and an explicit $\mathbf{P}$ hands every position its address from the first layer. GPT-2 learns $\mathbf{P}$ from scratch, the learned absolute positions of §5.4. Adding rather than concatenating keeps the width at $d$ , and training is free to give positions and tokens different directions in that space.

The learned table has a hard edge. There is no row for position $T$ , so a GPT-2 model cannot read token 1,025 at all. That's one reason later models switched to rotary embeddings (§6.7).

The block, with LayerNorm in front

Each of the $L$ blocks does two residual updates,

\mathbf{a} = \mathbf{x} + \operatorname{Attn}\!\big(\operatorname{LN}_1(\mathbf{x})\big), \qquad \mathbf{x}' = \mathbf{a} + \operatorname{MLP}\!\big(\operatorname{LN}_2(\mathbf{a})\big).

In plain words, each sublayer reads a normalized copy of the residual stream and adds its answer back to the unnormalized stream (§5.1). This is the pre-LN arrangement. GPT-1, like the original Transformer, normalized after each addition. GPT-2 moved the LayerNorm to the input of each sublayer and added one more LayerNorm after the last block (Radford et al., 2019). Pre-LN keeps a clean identity path from the embeddings to the output, which is why deep pre-LN stacks are far less sensitive to learning-rate warmup than post-LN ones (Xiong et al., 2020). GPT-2 reproductions still warm up, but a post-LN stack of the same depth often fails to train at all without it (§5.2 works through why).

LayerNorm normalizes each vector over its $d$ features with the population variance and a small $\epsilon = 10^{-5}$ , then applies a learned gain $\boldsymbol{\gamma}$ and bias $\boldsymbol{\beta}$ ,

\operatorname{LN}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \boldsymbol{\beta}, \qquad \mu = \frac{1}{d}\sum_j x_j, \quad \sigma^2 = \frac{1}{d}\sum_j (x_j - \mu)^2 .

Causal self-attention

The attention sublayer is the multi-head causal attention of §4.5. For one sequence, stack the normalized vectors into $\mathbf{X} \in \mathbb{R}^{T \times d}$ . GPT-2 computes queries, keys, and values with one fused matrix $\mathbf{W}_{QKV} \in \mathbb{R}^{d \times 3d}$ and splits the result,

[\mathbf{Q}\;\; \mathbf{K}\;\; \mathbf{V}] = \mathbf{X}\mathbf{W}_{QKV} + \mathbf{b}_{QKV}, \qquad \mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{T \times d}.

Each of the three is cut into $h$ heads of width $d_k$ (columns $i d_k$ to $(i+1)d_k$ belong to head $i$ ), and each head runs scaled dot-product attention with a causal mask $\mathbf{M}$ , whose entries are 0 where $j \le t$ and $-\infty$ where $j \gt t$ ,

\operatorname{head}_i = \operatorname{softmax}\!\left(\frac{\mathbf{Q}_i \mathbf{K}_i^\top}{\sqrt{d_k}} + \mathbf{M}\right)\mathbf{V}_i \in \mathbb{R}^{T \times d_k}, \qquad \operatorname{Attn}(\mathbf{X}) = [\operatorname{head}_1 \cdots \operatorname{head}_h]\,\mathbf{W}_O + \mathbf{b}_O .

In plain words, every position mixes the values of itself and the positions before it, weighted by how well its query matches their keys, in $h$ independent ways, and the output projection $\mathbf{W}_O \in \mathbb{R}^{d \times d}$ blends the heads back together. The scale is $\sqrt{d_k}$ , the head size, not $\sqrt{d}$ , for the variance reason of §4.3.

The MLP

The second sublayer works on each position separately, widening to $4d$ , applying GELU, and projecting back (§5.3),

\operatorname{MLP}(\mathbf{x}) = \operatorname{GELU}(\mathbf{x}\mathbf{W}_{fc} + \mathbf{b}_{fc})\,\mathbf{W}_{proj} + \mathbf{b}_{proj}, \qquad \mathbf{W}_{fc} \in \mathbb{R}^{d \times 4d},\; \mathbf{W}_{proj} \in \mathbb{R}^{4d \times d} .

GPT-2 uses the tanh approximation of GELU (Hendrycks & Gimpel, 2016),

\operatorname{GELU}(x) \approx \tfrac{1}{2}\, x \left(1 + \tanh\!\Big(\sqrt{2/\pi}\,\big(x + 0.044715\, x^3\big)\Big)\right),

which differs from the exact $x\,\Phi(x)$ by less than 0.001. That's small, but it matters when you load someone else's weights and expect identical outputs. PyTorch's F.gelu(x, approximate="tanh") is this version.

One ingredient is missing from the diagram, dropout. GPT-1 applied dropout with rate 0.1 to the embeddings, the attention weights, and each residual branch (Radford et al., 2018). Single-epoch pretraining on a large corpus rarely needs it, because the model sees each token about once and has little chance to memorize, so nanoGPT defaults to a rate of 0 for pretraining and suggests 0.1 or more for fine-tuning on small data. Our models leave it out to keep the code short. The widget model of later sections trains for many epochs on a small corpus, so it uses dropout. Whatever the rate, it must be switched off at evaluation time.

The head, and why it reuses the embedding

After the last block and the final LayerNorm, each position holds a vector $\mathbf{h}_t \in \mathbb{R}^d$ . The language-model head turns it into $V$ scores. GPT-2 does that with the same matrix it used at the bottom,

\mathbf{z}_t = \mathbf{E}\,\mathbf{h}_t \in \mathbb{R}^{V}, \qquad p(w_{t+1} = i \mid w_{\le t}) = \operatorname{softmax}(\mathbf{z}_t)_i = \frac{\exp(\mathbf{E}_{i,:}\cdot\mathbf{h}_t)}{\sum_{j}\exp(\mathbf{E}_{j,:}\cdot\mathbf{h}_t)} .

In plain words, the score of token $i$ is the dot product between the model's final vector and token $i$ 's own embedding. The model predicts a token by pointing its final vector toward that token's row of $\mathbf{E}$ .

This is weight tying (Press & Wolf, 2017; Inan et al., 2017), and §1.5 previewed it. There are two reasons to do it. The first is parameters. An untied head would be another $V \times d$ matrix, 38.6 million parameters for GPT-2 small, 31 percent of the model. The second is learning signal. With tying, a token's embedding gets gradient every time the token is predicted, and every time it is not predicted (through the softmax denominator), not only when it appears in the input. Press and Wolf found that tied models reach lower perplexity than untied ones of the same size.

The softmax-denominator gradient has a side effect worth knowing for interviews. A token that never occurs in the training data still gets pushed, on every step, toward "less likely than whatever came next". Its row drifts to an odd place in the embedding space, and if it ever shows up as input, the model sees a vector it never learned to read. That's one root of the under-trained "glitch tokens" of §2.6.

Tying is a trade-off, not a law. Many larger models untie the head, giving the input and output sides separate capacity, because at 7 billion parameters a second $V \times d$ matrix costs only a few percent (§6.7).

Interview

The gradient of a tied embedding

A tied embedding gets gradient from both ends of the model. Take this miniature model, with $\mathbf{E} \in \mathbb{R}^{V \times d}$ , $\mathbf{A} \in \mathbb{R}^{d \times d}$ , input ids idx of shape (n,), and targets of shape (n,),

\mathbf{X} = \mathbf{E}[\text{idx}], \qquad \mathbf{H} = \tanh(\mathbf{X}\mathbf{A}), \qquad \mathbf{Z} = \mathbf{H}\mathbf{E}^\top, \qquad \mathcal{L} = \frac{1}{n}\sum_{t} -\log \operatorname{softmax}(\mathbf{Z}_{t,:})_{y_t} .

Implement tied_lm_loss_and_grad(E, A, idx, targets) returning (loss, dE) with dE of shape (V, d). Both paths must be included, and idx may contain repeated ids.

import numpy as np

def tied_lm_loss_and_grad(E, A, idx, targets):
    """returns loss (float) and dE (V, d)"""
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Initialization, and the 1/√(2L) rule

GPT-2 initializes every weight matrix and both embedding tables from $\mathcal{N}(0, 0.02^2)$ , every bias to zero, and every LayerNorm to $\boldsymbol{\gamma} = \mathbf{1}$ , $\boldsymbol{\beta} = \mathbf{0}$ . The 0.02 is a fixed constant inherited from GPT-1, close to $1/\sqrt{d}$ for widths of a couple of thousand. (OpenAI's code used 0.01 for the position table, and reproductions commonly use 0.02 everywhere.) Then comes the one special rule. The two matrices that write into the residual stream, $\mathbf{W}_O$ and $\mathbf{W}_{proj}$ , are scaled down by $1/\sqrt{2L}$ , so their standard deviation is $0.02/\sqrt{2L}$ . The GPT-2 paper states it as scaling "the weights of residual layers at initialization by a factor of $1/\sqrt{N}$ where $N$ is the number of residual layers". The paper does not spell out $N$ , and reproductions such as nanoGPT read it as $2L$ , since each block has two residual branches.

The reason is the variance argument that §5.7 derives and measures. Each of the $2L$ residual branches adds a roughly independent random vector to the stream at initialization, variances add, and so the stream's variance grows linearly with depth unless each write is made $\sqrt{2L}$ times smaller. For GPT-2 small ( $L = 12$ ) that puts the two projections at $0.02/\sqrt{24} \approx 0.0041$ , and for GPT-2 XL ( $L = 48$ ) at $0.02/\sqrt{96} \approx 0.0020$ . In OpenAI's code both matrices are named c_proj, and in ours they are w_o and w_proj. Let's set up the model with exactly these values.

Python · cell 1

Our matrices are stored as (input, output), so a layer is x @ W, matching the math $\mathbf{X}\mathbf{W}$ . Now the forward pass, one function per box of the diagram.

Python · cell 2

Look at the head split, reshape(B, T, h, d_k).transpose(0, 2, 1, 3). Reshaping straight to (B, h, T, d_k) without the transpose also runs, and also produces the right shape, but it scrambles positions and features across heads. It's the most common silent bug in hand-written attention, and the exercise tests below catch it.

Now we can test the model the way you should test every new model, before training it.

Python · cell 3

Two checks passed. First, the loss at initialization is almost exactly $\ln V$ . Small random weights make every logit nearly equal, so the prediction is close to uniform, and a uniform guess over $V$ tokens costs $\ln V$ nats (§0.5). For GPT-2 that number is $\ln 50{,}257 = 10.82$ . An initial loss far above $\ln V$ means the initialization is too large and the model is confidently wrong before it has learned anything. Second, changing the token at position 20 leaves every earlier logit untouched to the last bit, which proves the mask is causal. A model that peeks at the future trains to a suspiciously low loss and then generates garbage.

Interview

The GPT forward pass, matched against a reference

Implement gpt_logits(P, idx, n_head) for a parameter dict shaped like init_gpt's, with idx of shape (B, T), returning logits of shape (B, T, V). It must be GPT-2 exactly. That means learned positions added to token embeddings, pre-LN blocks (LayerNorm with population variance and $\epsilon = 10^{-5}$ ), fused QKV split into n_head heads, attention scaled by $1/\sqrt{d_k}$ with a causal mask, tanh GELU, a final LayerNorm, and a head tied to wte. No loops over positions or heads. The tests compare against a reference to $10^{-8}$ and tell you which classic mistake you made if they disagree.

import numpy as np

def gpt_logits(P, idx, n_head):
    """returns (B, T, V) logits"""
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Counting every parameter

Now the puzzle from the top. Go through the diagram box by box. The embeddings hold $Vd$ token parameters and $Td$ position parameters. Each block holds

\underbrace{2d}_{\operatorname{LN}_1} + \underbrace{3d^2 + 3d}_{\mathbf{W}_{QKV},\,\mathbf{b}_{QKV}} + \underbrace{d^2 + d}_{\mathbf{W}_O,\,\mathbf{b}_O} + \underbrace{2d}_{\operatorname{LN}_2} + \underbrace{4d^2 + 4d}_{\mathbf{W}_{fc},\,\mathbf{b}_{fc}} + \underbrace{4d^2 + d}_{\mathbf{W}_{proj},\,\mathbf{b}_{proj}} = 12d^2 + 13d,

the final LayerNorm adds $2d$ , and the tied head adds nothing. So the total is

N_{\text{GPT-2}} = Vd + Td + L\,(12d^2 + 13d) + 2d .

In plain words, the familiar $12Ld^2$ counts only the blocks' matrices. A GPT also pays $Vd$ for its vocabulary, which for a small model is enormous.

For GPT-2 small, $Vd = 50{,}257 \times 768 = 38{,}597{,}376$ , $Td = 786{,}432$ , each block holds $12 \cdot 768^2 + 13 \cdot 768 = 7{,}087{,}872$ , the twelve blocks hold 85,054,464, and the final LayerNorm holds 1,536. The sum is 124,439,808. The 39 million "missing" parameters are the token embedding (38.6 million) and the position embedding (0.8 million), and without tying there would be another 38.6 million. The GPT-2 paper's table lists the four sizes as 117M, 345M, 762M, and 1542M, but the released checkpoints have the counts below, which is why everyone now says "124M".

Python · cell 4

Two conventions trip people up when comparing counts. Scaling-law papers usually report non-embedding parameters (Kaplan et al., 2020), everything except the token and position tables, which for GPT-2 small is $L(12d^2 + 13d) + 2d = 85{,}056{,}000$ , so the same model is "85M" in one paper and "124M" in another. And training repositories pad the vocabulary to 50,304, a multiple of 128 that makes the head's matmul faster (§6.7), which adds $47 \times 768 = 36{,}096$ rows that no token ever uses and gives 124,475,904. In memory the weights alone take 4 bytes each in float32, about 498 MB for GPT-2 small, or half that in bfloat16.

The share of parameters in the embedding shrinks from 31 percent to 5 percent as the model grows, because blocks grow as $d^2$ and the vocabulary only as $d$ . That asymmetry is behind several decisions in §6.7, including untying the head in big models. Explore it below, including the GPT-3 configuration (Brown et al., 2020) and our tiny model.

Core

Count a GPT exactly

Write count_gpt2_params(V, T, d, L, tied=True, bias=True) returning the exact parameter count of the GPT-2 architecture. With tied=False the head is a separate $V \times d$ matrix (still without a bias). With bias=False every linear layer and every LayerNorm has no bias (LayerNorms keep their gains), as in models that drop biases.

def count_gpt2_params(V, T, d, L, tied=True, bias=True):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

In PyTorch

Here is the same model as a PyTorch module, written the way nanoGPT and the "Let's reproduce GPT-2" video do it. F.scaled_dot_product_attention with is_causal=True replaces our masked softmax and dispatches to a fused kernel (FlashAttention, §12.3) on a GPU. The meta device builds the full GPT-2 small with shapes but no memory, so we can count its parameters instantly.

🔥 In PyTorchfor a machine with PyTorch

model.parameters() yields each tensor once, so the tied matrix is counted once, which is what we want. (Summing over state_dict() instead would count it twice, since the state dict lists both names.) Now the real test. Copy our NumPy parameters into the PyTorch model and compare logits in float64.

🔥 In PyTorchfor a machine with PyTorch

They agree to about $10^{-16}$ , which is float64 rounding. Notice the transposes. nn.Linear stores its weight as (output, input) and computes $\mathbf{x}\mathbf{W}^\top$ , while our NumPy matrices, like the Conv1D layers in OpenAI's GPT-2 checkpoint, are stored as (input, output). Forgetting that transpose is the classic bug when loading GPT-2 weights into your own module. With square matrices like $\mathbf{W}_O$ it doesn't even raise a shape error.

design choice	GPT-1 (2018)	GPT-2 (2019)	LLaMA-style (§6.7)
normalization	post-LN, LayerNorm	pre-LN plus a final LayerNorm	pre-norm RMSNorm
positions	learned absolute, 512	learned absolute, 1,024	rotary (RoPE), none in the input
MLP	GELU, $4d$	GELU (tanh), $4d$	SwiGLU, about $\tfrac{8}{3}d$
biases	yes	yes	none
head	tied	tied	usually untied
residual init	$\mathcal{N}(0, 0.02^2)$	scaled by $1/\sqrt{2L}$	varies by codebase

🎯 Interview questions on this section

Walk me through GPT-2's forward pass with shapes.

Ids $(B, T)$ index the token table and add the learned position rows, giving $(B, T, d)$ . Each of $L$ pre-LN blocks computes $\mathbf{x} + \operatorname{Attn}(\operatorname{LN}(\mathbf{x}))$ , with a fused $(d, 3d)$ QKV projection split into $h$ heads of size $d/h$ , causal scaled dot-product attention, and a $(d, d)$ output projection, then $\mathbf{x} + \operatorname{MLP}(\operatorname{LN}(\mathbf{x}))$ with $d \to 4d \to d$ and tanh GELU. A final LayerNorm and the transposed token table give logits $(B, T, V)$ . The logits at position $t$ predict token $t+1$ .

How many parameters does GPT-2 small have, and where are they?

124,439,808. Each block has $12d^2 + 13d = 7{,}087{,}872$ , so the 12 blocks hold 85.1M. The token embedding holds $50{,}257 \times 768 = 38.6$ M, the position embedding 0.79M, and the final LayerNorm 1,536. The head is tied, so it adds nothing. Untied, it would add another 38.6M.

What is weight tying, and what are its pros and cons?

The output projection reuses the input embedding matrix, so logits are $\mathbf{E}\mathbf{h}$ . It saves $Vd$ parameters, a large fraction for small models, and every embedding row gets gradient on every step through the softmax, which Press and Wolf found lowers perplexity. The cost is that input and output representations are forced to be the same, and that rare or never-seen tokens get only the push-down gradient from the softmax, leaving odd vectors that behave strangely as inputs. Large models often untie because the extra $Vd$ is a small fraction of their size.

Why does GPT-2 scale the residual projections by 1/sqrt(2L) at initialization?

The final residual stream is the embedding plus $2L$ sublayer outputs. At initialization those outputs are roughly uncorrelated with similar variance, so the stream's variance grows linearly with depth. Scaling each branch's last matrix by $1/\sqrt{2L}$ divides each contribution's variance by $2L$ , so the total stays roughly constant as the model gets deeper, keeping activations and gradients at a depth-independent scale.

What loss should a freshly initialized GPT have, and what does it mean if it's much higher?

About $\ln V$ , since tiny random logits give a nearly uniform prediction. That's 10.82 for GPT-2's vocabulary and 4.03 for our 56 characters. A much higher value means the logits are already large, the model is confidently wrong at step zero, which points to an initialization or scaling bug, such as a missing $1/\sqrt{d_k}$ , a large init on the head, or an untied head initialized with a large standard deviation.

Why pre-LN instead of post-LN?

Post-LN normalizes after each residual addition, so the gradient to early layers passes through every LayerNorm, and deep post-LN models need careful learning-rate warmup to avoid diverging. Pre-LN leaves an unnormalized identity path from the embedding to the final LayerNorm, so gradients reach every layer directly, and deep stacks train stably. The cost is that the stream can grow with depth, which is why GPT-2 adds a final LayerNorm and scales the residual initialization.

🧭 What to remember

GPT-2 is embeddings (tokens plus learned positions), $L$ pre-LN blocks of causal multi-head attention and a $4d$ GELU MLP, a final LayerNorm, and a head tied to the token embedding.
The exact count is $Vd + Td + L(12d^2 + 13d) + 2d$ , which gives 124,439,808 for GPT-2 small. The vocabulary is 31 percent of a small model and a few percent of a big one.
Tying saves $Vd$ parameters and gives every embedding row a gradient on every step, including a push-down for tokens that never occur.
Initialize with $\mathcal{N}(0, 0.02^2)$ and shrink the residual projections by $1/\sqrt{2L}$ , so the residual stream's scale does not grow with depth.
Test every new model before training. The loss at init should be $\ln V$ , perturbing a future token must not change past logits, and two independent implementations should agree to rounding error.
Next, §6.3 trains this model and teaches you to read what the loss curve is telling you.

Radford et al. (2018). Improving Language Understanding by Generative Pre-Training. GPT-1, the 12-layer decoder, 0.02 initialization, and the pretrain-then-fine-tune recipe.
Radford et al. (2019). Language Models are Unsupervised Multitask Learners. GPT-2's changes in §2.3, including pre-LN, the final LayerNorm, and the residual initialization.
Press & Wolf (2017). Using the Output Embedding to Improve Language Models. Why tying the input and output embeddings helps.
Hendrycks & Gimpel (2016). Gaussian Error Linear Units (GELUs). The activation and its tanh approximation.
Karpathy. nanoGPT model.py. A compact reference implementation of exactly this model, including loading OpenAI's weights.
Xiong et al. (2020). On Layer Normalization in the Transformer Architecture. Why pre-LN trains with far less dependence on warmup than post-LN.
Haviv et al. (2022). Transformer Language Models without Positional Encodings Still Learn Positional Information. How much order a causal mask leaks on its own.
Kaplan et al. (2020). Scaling Laws for Neural Language Models. Its scaling laws count non-embedding parameters, the convention behind "85M" for GPT-2 small.

Zero to Frontier is made for larger screens

The model

The whole model in one picture

Embeddings, tokens plus positions

The block, with LayerNorm in front

Causal self-attention

The MLP

The head, and why it reuses the embedding

The gradient of a tied embedding

Initialization, and the 1/√(2L) rule

The GPT forward pass, matched against a reference

Counting every parameter

Count a GPT exactly

In PyTorch

Read more

Watch next