The GPT architecture

One table, two jobs

Your phone's contacts list works in two directions. Type a name and it finds the number. When a call comes in, it matches the number back to a name. One list does both jobs.

Many GPTs do the same with the token embedding table. On the way in, the table turns a token into a vector. On the way out, the LM head must turn a vector back into a score for each token. It can reuse the same table. It takes the dot product of the final vector with every token's row. Rows that match the vector well get high scores. Sharing the table like this is called weight tying.

The table has one row per token and one column per number in the vector. Write $V$ for the vocabulary size. With width $d$, the table is $V \times d$. A parameter is one learned number in the model. Tying stores the table once instead of twice, which saves $V \cdot d$ parameters.

For example, with 200 tokens and width 64, the table holds 200 × 64 = 12,800 numbers. Tying saves all 12,800 of them.

A complete small GPT

Here is a whole GPT you can run. The Block class is a compact transformer block. The GPT class follows the four steps. The line marked # weight tying makes the head share the token table.

import math
import torch
import torch.nn as nn

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=-1)
        q, k, v = (t.view(B, T, self.h, d // self.h).transpose(1, 2) for t in (q, k, v))
        s = (q @ k.transpose(-2, -1) / math.sqrt(d // self.h)).masked_fill(~torch.tril(torch.ones(T, T)).bool(), float("-inf"))
        x = x + self.proj((torch.softmax(s, dim=-1) @ v).transpose(1, 2).reshape(B, T, d))
        return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self, vocab, context, d, n_layers, n_heads):
        super().__init__()
        self.context = context
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(context, d)
        self.blocks = nn.Sequential(*[Block(d, n_heads) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok.weight                  # weight tying

    def forward(self, idx):
        T = idx.shape[1]
        x = self.tok(idx) + self.pos(torch.arange(T))
        return self.head(self.ln_f(self.blocks(x)))

torch.manual_seed(0)
model = GPT(vocab=200, context=32, d=64, n_layers=2, n_heads=4)
idx = torch.randint(0, 200, (3, 10))
print("logits:", tuple(model(idx).shape), "  parameters:", f"{sum(p.numel() for p in model.parameters()):,}")

It prints (3, 10, 200) for the logits: 3 sequences, 10 tokens each, 200 scores per token. It also prints 114,944 parameters. Most of them sit in the two blocks, at 49,984 each. The token table holds the 12,800 you worked out above.

Here torch is the course's small in-browser version of PyTorch. The next lesson trains this exact model.

Counting the parameters

A model's size is usually given as its number of parameters. You can work it out from four numbers on the model card. Here are the symbols, one per line:

• $V$ is the vocabulary size.
• $n_{\text{ctx}}$ is the context length, the longest sequence the model can take. It's also the number of rows in the position table.
• $d$ is the width.
• $L$ is the number of layers, which means the number of blocks.

In plain words: count the token table, $V$ rows of $d$ numbers. Add the position table, $n_{\text{ctx}}$ rows of $d$ numbers. Add one block's worth for every layer. Add $2d$ for the final layer norm. The head adds nothing, because it's tied.

The $12 d^2 + 13 d$ is one block's count. The transformer block lesson showed that a block holds about $12 d^2$ parameters. Attention has $4 d^2$ of them and the MLP has $8 d^2$. The $13 d$ covers small extras: the bias terms and the layer norms. The panel below has the full tally.

The number of heads isn't in the formula. As the multi-head lesson showed, splitting into heads adds no parameters.

V, n_ctx, d, L = 10, 4, 2, 1
tokens = V * d                          # token table
positions = n_ctx * d                   # position table
blocks = L * (12 * d * d + 13 * d)      # every block
final_norm = 2 * d                      # final layer norm (the head is tied)
print(tokens, positions, blocks, final_norm)
print("total:", tokens + positions + blocks + final_norm)

Counting GPT-2 small

The same formula works for real models. GPT-2 small has a vocabulary of 50,257 and a context of 1,024. Its width is 768, with 12 layers and 12 heads. Its head is tied. The function below is the formula, with the block split into its parts. It also counts the three bigger GPT-2 models.

def gpt_params(vocab, context, d, n_layers):
    embeddings = vocab * d + context * d
    per_block = 4 * d * d + 4 * d          # attention: Q, K, V, output projection + biases
    per_block += 8 * d * d + 5 * d         # MLP: d → 4d → d, with biases
    per_block += 4 * d                     # two layer norms (scale and shift)
    return embeddings + n_layers * per_block + 2 * d   # + final layer norm (head is tied)

configs = {"GPT-2 small": (50257, 1024, 768, 12), "GPT-2 medium": (50257, 1024, 1024, 24),
           "GPT-2 large": (50257, 1024, 1280, 36), "GPT-2 XL": (50257, 1024, 1600, 48)}
for name, cfg in configs.items():
    print(f"{name:<13} {gpt_params(*cfg) / 1e6:7.1f} million parameters")

GPT-2 small comes out at 124.4 million, its published size. Medium, large and XL come out at about 355 million, 774 million and 1.56 billion. Those match the published family. About 39 million of small's parameters are the token table alone.

table = 50257 * 768
total = 124439808
print(round(table / total * 100))

Generation: predict, sample, append, repeat

Think of the word suggestions on your phone's keyboard. Tap one, and it suggests the next word based on everything typed so far. A GPT writes text the same way, one token at a time:

1. Run the model on the tokens so far.
2. Take the logits at the last position.
3. Pick one token, using a sampling strategy from the sampling lesson.
4. Add that token to the end of the sequence.
5. Repeat.

@torch.no_grad()
def generate(model, idx, new_tokens, temperature=1.0):
    for _ in range(new_tokens):
        logits = model(idx[:, -model.context:])[:, -1, :]       # last position only
        probs = torch.softmax(logits / temperature, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, next_id], dim=1)
    return idx

Here's what each line does.

• @torch.no_grad() tells PyTorch not to record gradients. Nothing is being trained here, so that work would be wasted.
• idx[:, -model.context:] keeps only the last model.context tokens. The model has one learned position vector per slot, up to its context length. It has none for later slots, so longer input must be cut down first.
• [:, -1, :] picks the logits at the last position.
• Dividing by temperature before the softmax sharpens or flattens the probabilities, as in the sampling lesson.
• torch.multinomial draws one token at random, weighted by the probabilities.
• torch.cat adds the new token to the end.

This simple loop redoes the work for the whole sequence at every step. Real systems save the keys and values from earlier steps instead. That store is the KV cache, which has its own lesson.

Reading a configuration

A model card lists a few settings, called its configuration. Here's what each one controls, with GPT-2 small's values.

Width is often written $d_{\text{model}}$ ("d model"). It's the same $d$ as in the formula.

The biggest models today scale up every row. Their widths run into the thousands. They have around a hundred layers and vocabularies of 100,000 tokens or more. Their context windows reach hundreds of thousands of tokens, sometimes more. The skeleton is still the one in this lesson.