Building a complete GPT

The last two lessons built the parts: attention lets each character gather information from the ones before it, multiple heads let it gather several kinds at once, and position embeddings let it know the order. This lesson assembles the parts into a complete GPT.

The full code is in code/09-transformer/gpt.py, under 150 lines. Its structure is the same as GPT-2's, only much smaller. The training and generation in the next two lessons both use it.

python inspect_gpt.py

Overall structure

  词元编号
     │
  词元嵌入 + 位置嵌入
     │
  ┌──────────────────────────┐
  │  LayerNorm → 多头注意力   │ ─┐
  │        + ←────────────────── ┘ 残差
  │  LayerNorm → 前馈网络     │ ─┐
  │        + ←────────────────── ┘ 残差
  └──────────────────────────┘
     │   （这样的块叠 4 层）
  LayerNorm
     │
  输出层：对词表里每个字打分

Reading from the bottom up: each character first becomes a vector (token embedding plus position embedding), passes through several identically structured Transformer blocks, and finally the output layer scores the next character. We've already written the attention inside a block; there are three new things too: a feed-forward network, residual connections and LayerNorm.

The feed-forward network: each character thinks for itself

Attention passes information between characters. After gathering information, each character also needs to "digest" it on its own, and that's the feed-forward network:

self.mlp = nn.Sequential(  # 前馈网络：先放大 4 倍，过激活函数，再缩回来
    nn.Linear(cfg.n_embd, 4 * cfg.n_embd),
    nn.GELU(),
    nn.Linear(4 * cfg.n_embd, cfg.n_embd),
    nn.Dropout(cfg.dropout),
)

It's just the most ordinary kind of two-layer network from Module 08: expand 128 dimensions to 512, apply an activation function, and shrink back to 128. GELU is an activation function much like ReLU, except it's a smooth curve near 0; GPT-2 uses it.

The feed-forward network computes each position separately, with no influence between positions. So the division of labour in a Transformer block is: attention handles "communicating", the feed-forward network handles "thinking". As you'll see below, most of the model's parameters are actually in the feed-forward networks.

Residual connections: change the original a little

def forward(self, x, cache=None):
    x = x + self.attn(self.ln1(x), cache)  # 残差连接：在原来的基础上加一点修改
    x = x + self.mlp(self.ln2(x))
    return x

Note that this is x = x + ..., not x = self.attn(x). What attention and the feed-forward network compute isn't a new vector but an "amount of change" to the original, added back onto it. This is called a residual connection.

Why do it this way? As Module 08, Lesson 3 explained, gradients are multiplied back layer by layer in backpropagation. With many layers, gradients multiplied many times easily become tiny (nothing is learned) or huge (training blows up). With residual connections, the derivative of x + f(x) with respect to x always contains a 1, so gradients can travel straight back to earlier layers along the "plus" path without passing through each layer's transformation. This is one of the keys to training deep networks of dozens or hundreds of layers.

Another way to see it: the whole model has a "main road" running straight from input to output, and each block just adds a little to it.

LayerNorm: keeping values stable

Module 08, Lesson 2 showed that when inputs are on very different scales, training is hard. As computation passes layer after layer, the numbers in the vectors also grow larger and larger or smaller and smaller. Before each sublayer, LayerNorm adjusts each vector to mean 0 and standard deviation 1, then multiplies and adds two sets of learnable parameters (letting the model decide the right range itself).

It's the same idea as Module 08's standardisation, except it's applied to each position's vector separately, and repeatedly in the middle of the network.

Note that LayerNorm goes before attention and the feed-forward network (self.attn(self.ln1(x))), not after. This is called Pre-LN; models since GPT-2 have almost all placed it this way, since it trains more stably. There's one more LayerNorm (ln_f) just before the output.

Input and output

self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)
...
self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
self.head.weight = self.tok_emb.weight  # 输出层和词元嵌入共用一套参数，GPT-2 也是这样做的

nn.Embedding is just a table: row i is the vector for character i. The token embedding table has 6,289 rows (the vocabulary size), and the position embedding table has 128 rows (handling at most 128 positions).

The output layer turns the 128-dimensional vector into 6,289 scores, one per character. There's a small trick here: the output layer and the token embedding share the same matrix. The token embedding turns "character" into "vector", and the output layer does the reverse, judging which "character" a "vector" is most like, so using the same parameters is reasonable, and it saves a large chunk of parameters.

The forward pass strings all of this together:

def forward(self, idx, targets=None, caches=None, start=0):
    B, T = idx.shape
    pos = torch.arange(start, start + T, device=idx.device)
    x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
    for i, block in enumerate(self.blocks):
        x = block(x, None if caches is None else caches[i])
    logits = self.head(self.ln_f(x))  # (B, T, 词表大小)：每个位置对下一个词元的打分
    loss = None
    if targets is not None:
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
    return logits, loss

caches and start are for Lesson 6's KV cache; ignore them for now. The loss is Module 08, Lesson 5's cross-entropy: every position is a classification problem of "choose one of 6,289 characters".

Shapes

== 2. 一批数据流过模型时的形状
  输入的词元编号        (2, 6)
  嵌入之后              (2, 6, 128)
  经过 4 个块之后        (2, 6, 128)
  输出层                (2, 6, 6289)：每个位置对 6289 个字各打一个分

Two lines of poetry, 6 characters each, go in as (2, 6). After embedding, each character becomes a 128-dimensional vector. Through 4 blocks, the shape doesn't change at all, which is exactly why blocks can be stacked as deep as you like. Finally each position gets 6,289 scores.

Note that there's output at every position: a 6-character text provides 6 "guess the next character" questions at once during training. Position 1 sees "白" and guesses "日", position 2 sees "白日" and guesses "依"... Because of the causal mask, none of the questions can peek at its answer. This makes training much more efficient.

Where the parameters are

== 1. 参数都在哪里（词表 6289 个字符）
  词元嵌入   6289 × 128 = 804,992（输出层和它共用，不另算）
  位置嵌入   128 × 128 = 16,384
  每个块     198,272：注意力 66,048，前馈网络 131,712，两个 LayerNorm 512
  4 个块共  793,088
  合计       1,614,720，其中词元嵌入占 50%

1.61 million parameters in all. You can work it out yourself:

• Attention: qkv is 128×384 plus 384 biases, proj is 128×128 plus 128, 66,048 in total.
• Feed-forward network: 128×512 plus 512, 512×128 plus 128, 131,712 in total, twice attention.
• LayerNorm: 128 scales and 128 offsets each, 512 for the two.

In our small model, half the parameters are in the token embedding table, because the vocabulary has 6,289 characters while the vectors are only 128-dimensional. LLMs have a completely different ratio: vector dimensions in the thousands, dozens of layers, and far more parameters in the blocks than in the embeddings. Within each block, the feed-forward network takes about two thirds and attention about one third, a ratio that's roughly the same in LLMs.

Before training

== 3. 没训练过的模型，损失应该接近随便猜
  损失 8.775，ln(6289) = 8.747
  下一个字最可能是：骠蔫潮涨湄（随机初始化，毫无道理）

An untrained model scores all 6,289 characters almost equally, which amounts to guessing at random, so the cross-entropy should be ln(6289) = 8.747. It's actually 8.775, very close. As Module 08, Lesson 5 said: if the loss at the start of training is far from this number, something is probably wrong. It's the first check to do once a model is built.

The model's initial parameter values are drawn from a normal distribution with std=0.02 (the _init method), so at the start every character's score is about the same and the loss is close to random guessing.

How GPT differs from real LLMs

This GPT's structure is essentially the same as GPT-2's. As of September 2026, mainstream open-source LLMs have made some improvements on this skeleton; common ones include:

• Position information using the rotary position embedding (RoPE) from last lesson.
• LayerNorm replaced by the simpler-to-compute RMSNorm.
• The feed-forward network replaced by a "gated" structure (such as SwiGLU).
• Several attention heads sharing K and V (grouped-query attention), reducing the memory the KV cache takes in the next lesson.
• The feed-forward network replaced by many "experts", with each character using only a few of them (mixture of experts, MoE); DeepSeek's models have this structure.

But the skeleton is unchanged: embeddings, several "attention + feed-forward network" blocks, residual connections, normalisation, an output layer. Once you understand this 150-line GPT, reading those models' code, you'll recognise most of it.

Exercises

1. Change n_layer to 8 and n_embd to 256. How many parameters are there now, and what share is the token embedding? Work it out with this lesson's formulas first, then run it to check.
2. Remove the two residual connections in Block.forward (x = self.attn(self.ln1(x))), train for 300 steps with next lesson's train.py, and compare with the original loss.
3. Delete the line that shares parameters between the output layer and the token embedding. How many more parameters are there?

Self-check

1. What parts does a Transformer block have? How do attention and the feed-forward network divide the work?

Two sublayers, multi-head causal self-attention and a feed-forward network, each preceded by a LayerNorm and wrapped in a residual connection. Attention passes information between positions, and the feed-forward network computes on each position separately, processing the information gathered.

2. Why do residual connections help train very deep networks?

A residual connection adds each sublayer's output back onto its input, so the derivative of x + f(x) with respect to x always contains a 1. In backpropagation, gradients can travel back to earlier layers along this direct path, without becoming tiny or huge from being multiplied through many layers.

3. Why should an untrained model's loss be close to ln(vocabulary size)?

Randomly initialised parameters are small, so the model scores every character about the same, and after softmax each character's probability is about 1/vocabulary size. Cross-entropy is -log(probability of the correct character), that is, -log(1/vocabulary size) = ln(vocabulary size).