Part III — The GPT

Assembling GPT & generating text

Lesson 07 built one transformer block: a self-contained unit that maps a tensor of shape (B, T, d) back to the same shape, so it can be stacked. But a single block is shallow, and its output is still a stream of vectors — not text. This lesson stacks L of those blocks, caps them with a final normalization and a head that projects back to the vocabulary, and then runs the autoregressive loop from lesson 00 over the real full-model shapes to emit actual tokens.

1 · The full stack, end to end

Everything in the last seven lessons was a component. Now we bolt them into one module. Reading bottom to top, a GPT is five stages, and it is worth tracing the tensor shape through each so nothing feels like magic.

Start with a batch of token ids, shape (B, T) — B sequences, each T tokens long. The token embedding table (lesson 02) looks up a d-vector per id, and the positional embedding table adds a per-slot vector so the model can tell position 0 from position 5; their sum, after a dropout, is the input to the stack: shape (B, T, d). This is the residual stream — the width-d bus that every block reads from and writes back to.

Then the heart: L transformer blocks in sequence, each the pre-norm causal-attention-plus-MLP unit from lesson 07. Because a block preserves shape, the output after block 12 is still (B, T, d) — but every vector has now been rewritten L times, absorbing context from the whole (masked) sequence at every layer. A final LayerNorm standardizes those vectors one last time so the head sees a well-scaled input.

Finally the LM head: a single linear map from d to V, no bias. It takes each position's d-vector and produces V numbers — one raw score (a logit) per vocabulary token. The output is (B, T, V): for every position in every sequence, a full slate of next-token scores. That is the type signature from lesson 00 made concrete.

class GPTModel:
    tok_emb   : Embedding(V, d)          # id  -> vector
    pos_emb   : Embedding(context_len, d)# slot-> vector
    blocks    : [TransformerBlock] * L   # the stack (lesson 07)
    final_ln  : LayerNorm(d)
    lm_head   : Linear(d, V, bias=False) # vector -> V logits

    forward(ids):                        # ids : (B, T)
        x = tok_emb(ids) + pos_emb(range(T))   # (B, T, d)
        x = dropout(x)
        for block in blocks: x = block(x)       # (B, T, d)
        x = final_ln(x)
        return lm_head(x)                       # (B, T, V)  logits

2 · Counting the parameters

The stack looks intricate, but its parameter count collapses to two terms, and the derivation is worth doing once because it explains where a GPT's mass actually lives. Almost everything is in the L repeated blocks; the embeddings and head are a thin skin on top.

Inside one block, ignore the small LayerNorm scale/shift vectors (they contribute only O(d) each). Two big pieces remain:

• Attention: four d×d matrices — W_q, W_k, W_v and the output projection W_o — for 4d² parameters.
• Feed-forward MLP: two matrices, d→4d and 4d→d, for 4d² + 4d² = 8d² parameters.

So one block holds about 4d² + 8d² = 12d² parameters, and L of them give 12Ld². Outside the stack, the dominant term is the vocabulary-sized embedding (or head), V·d. Positional embeddings add only context\_length·d, small by comparison. The whole model is therefore, to leading order,

Now plug in GPT_CONFIG_124M — the canonical GPT-2-small configuration: V = 50,257, context\_length = 1,024, d = 768, h = 12 heads, L = 12 layers, drop_rate 0.1, no QKV bias (so d\_head = 768/12 = 64):

That is the "124M" in the name. If you keep a separate LM head, you pay the V·d term twice — once for the embedding, once for the head — which is why an actual PyTorch numel() on the untied model reports ~163M, and subtracting the head's 38.6M lands you back at the tied 124M. At float32 (4 bytes each) that untied model is about 622 MB of weights.

3 · From logits to text

The model emits logits, but a person wants words. Bridging that gap is the generation loop — and it is exactly the loop you already drove in lesson 00, now running over the real (B, T, V) shapes instead of a toy score table.

Given a prompt of ids, do one forward pass. The output is (B, T, V), but only the last position matters: the logits at slot T-1 are the model's prediction for the token that comes after the prompt (this is why the causal mask from lesson 06 was essential — position T-1 has legitimately seen all of 0..T-1 and nothing after). Slice that row out — shape (B, V) — turn it into a probability distribution with softmax, and pick a token. The simplest rule is greedy decoding: take the argmax, the single highest-scoring token.

generate(model, ids, n_new, context_len):     # ids : (B, T)
    for _ in range(n_new):
        ids_cond = ids[:, -context_len:]       # crop to the last context_len tokens
        logits   = model(ids_cond)             # (B, T', V)
        logits   = logits[:, -1, :]            # last position only -> (B, V)
        probs    = softmax(logits, dim=-1)     # (B, V)
        next_id  = argmax(probs, dim=-1)       # greedy -> (B, 1)
        ids      = concat([ids, next_id], dim=1)  # append, sequence grows by 1
    return ids

Two subtleties earn their keep. First, softmax before argmax is technically redundant for greedy decoding — softmax is monotonic, so the argmax of the probabilities equals the argmax of the raw logits. We keep it because it names the right mental model (logits → distribution → sample) and because lesson 11 will replace argmax with real sampling, where the probabilities genuinely matter. Second, and non-negotiable: crop to the context length. The model has a fixed positional table of size context\_length (1,024 for GPT-2); feed it a sequence longer than that and there is no position embedding for the overflow and the forward pass breaks. So each step keeps only the last context\_length tokens as input. The generated text can be arbitrarily long; the window the model conditions on is bounded.

That is the whole generation algorithm: forward, slice the last position, pick, append, crop, repeat. Wrap it in a loop with a token budget and you have generate_text_simple.

4 · Why raw output is gibberish

Assemble the model, initialize it, hand it the prompt "Hello, I am" and generate a few tokens. Out comes something like "Hello, I am Featureiman Byeswickattribute argue" — grammatically shaped noise. Nothing is broken. The architecture is exactly GPT-2's, the shapes all check out, the loop runs. The problem is upstream of all of it: the weights are random.

Every matrix — the embeddings, all 12·12 attention and MLP matrices, the head — was filled with small random numbers at initialization. So the logits are essentially random, the argmax picks a near-arbitrary token, and the loop faithfully strings arbitrary tokens together. The model has a perfectly good mechanism for turning context into a next-token distribution; it has simply never been told which distributions are good. It has no notion of "wrong."

This is the wall that ends Part III and opens Part IV. We have a machine with the right shape and 124 million knobs, and no signal telling us how to turn them. To train it, we need a single number that says how far the model's predicted distribution is from the truth at each position — a number that is small when the model is confident and correct and large when it is confident and wrong. That number is the cross-entropy loss, and defining it is the entire job of lesson 09.

5 · Assemble it, then watch it generate

The widget has two halves that together prove the whole lesson. On the left, an assembler: drag d and L (heads and vocabulary are pinned to the GPT-2 config) and watch the parameter count N = 12Ld² + Vd update live, with a bar splitting the total into its "blocks" and "vocab" parts — proving where a GPT's mass lives. On the right, the generation loop from lesson 00, now framed as running over these full-model shapes: Step does one forward-pick-append, Run loops it, and the context-window slider crops the input so you can watch the loop conditioning only on the last few tokens even as the generated sequence grows past the window. The point: a GPT is embeddings + L blocks + a head, and generation is nothing more than lesson 00's loop over real shapes.

Failure modes & checklist

Checkpoint

Where this points next

You have now assembled the entire architecture — stack, norm, head — and watched it generate text by the same loop that opened the series. But the text is nonsense, and nothing in the model can tell it so. The forward pass produces a distribution; we have no way to say how good that distribution is against the actual next token in the data. Before we can train a single weight, we need a scalar objective: a number that is large when the model puts low probability on the true next token and small when it puts high probability there. That number is cross-entropy, and its exponential is perplexity — the subject of lesson 09, The loss — cross-entropy & perplexity.

Interview prompts

• Trace the tensor shape from token ids to logits through a GPT. (§1 — (B,T) ids → (B,T,d) after token+pos embeddings → (B,T,d) through the L blocks and final LayerNorm → (B,T,V) after the LM head.)
• Derive the parameter count N ≈ 12Ld² + Vd. (§2 — per block: attention 4d² (W_q,W_k,W_v,W_o) + MLP 8d² (d→4d→d) = 12d²; times L blocks, plus the vocabulary embedding/head Vd.)
• Why is GPT-2-small quoted at 124M when a PyTorch count gives ~163M? (§2 — weight tying reuses the token embedding as the LM head, charging the Vd ≈ 39M term once; the untied model pays it twice.)
• In generation, why do you use only the last position's logits? (§3 — with the causal mask, position T-1 is the one that has attended to the whole prompt, so its logits predict the token after the prompt; earlier rows predict earlier tokens.)
• Why must you crop the input to the context length each step? (§3 — the positional embedding table has only context\_length rows; a longer input has no position vector for the overflow and the forward pass breaks. The generated text can exceed it; the conditioning window cannot.)
• Is the softmax before argmax necessary for greedy decoding? (§3 — no; softmax is monotonic so argmax of probabilities equals argmax of logits. It is kept for clarity and because lesson 11's sampling needs the actual probabilities.)
• Why does a freshly assembled GPT output gibberish, and what does that motivate? (§4 — the architecture is correct but the weights are random, so the logits are near-arbitrary; the model has no notion of "wrong," which motivates the cross-entropy loss of lesson 09.)

Companion reads: cs336 · 02 The Transformer counts the same parameters and the matching FLOPs at scale; gpt_mini · 01 Architecture shows this assembled stack as runnable code.