CapstonePython

GPT Model Assembly

A GPT model is twelve stacked blocks, a token embedding, a learned position embedding, a final LayerNorm, and a language model head tied to the token embedding. Assemble them, count the parameters against the 124 million reference, then generate text with multinomial sampling, temperature, and top-k.

TIME: ~90 minutesPREREQUISITES: Phase 19 lessons 30 to 34

Learning Objectives

Assemble the transformer block from lesson 34 into a full GPT model: token embedding, position embedding, N blocks, final LayerNorm, language model head.
Reproduce the 124 million parameter configuration: vocab 50257, context 1024, embedding 768, twelve heads, twelve layers.
Tie the language model head weights to the token embedding and explain why that saves ~38 million parameters at this scale.
Generate text from a prompt with multinomial sampling, temperature scaling, and top-k truncation, holding context length with a sliding window.
Measure parameter count and forward pass cost against the 124M target.

The Problem

A transformer block maps vectors to vectors and nothing else. To make a language model you wrap it in four steps: turn token ids into vectors, mix in position information, run the result through the stack, and project back to vocabulary logits. Skip one and the model fails to forward, loses track of word order, or has no way to emit a token.

The shape of the model matters too. GPT-2 small has 124 million parameters at exactly the configuration above, and the number falls out of the parts. Vocab 50257 times embedding 768 is the token table. Position 1024 times 768 is the position table. Twelve blocks at roughly 7 million parameters each come to 84 million. The final head reuses the token table through weight tying, so it adds nothing. Sum the pieces and you land on 124 million. If your parameter count misses the reference, you wired something wrong.

Warm-up0/3

1. Which PyTorch layer turns integer token ids into learned vectors by looking up rows of a table?Warm-up

2. A transformer block takes a tensor of shape (batch, seq, d_model). What shape does it return?Warm-up

3. Attention scores depend only on which vectors are present, not their order. What does a model need to add so word order matters?Warm-up

Check: The Problem0/5

1. How do the 124 million parameters of the reference configuration break down?Explain

2. What goes wrong if the step that projects the final vectors to vocabulary logits is left out?Explain

3. You swap the learned position embedding for a sinusoidal table computed at construction. By how much does the 124M parameter count drop?Apply

4. You build the tied 124M model with unique-parameter counting and get about 163 million. What is the most likely cause?Analyze

5. A new GPT assembly forwards and emits tokens, but the model clearly cannot tell word order. Which skipped step explains it?Analyze

The Concept

Mermaid diagram

A GPT model is a token embedding, a position embedding, a stack of transformer blocks, a final LayerNorm, and a language model head that maps back to vocabulary logits. Token ids become token vectors and position ids become position vectors. The two are added and sent through the stack. The final LayerNorm is the one piece outside the blocks that survives every modern variant, and the LM head reuses the token embedding matrix, which is what weight tying means.

Weight tying

Weight tying means the LM head and the token embedding are the same parameter tensor, used twice. The embedding has shape (vocab, d_model), and the head has to project from d_model back to vocab, which is the transpose of that shape. At vocab 50257 and d_model 768 the matrix holds 38 million parameters. Untied, you pay for it twice. Tied, you pay once, and you also get a slightly cleaner gradient signal because the embedding and the head update together.

Position embedding is learned, not sinusoidal

A learned position embedding is a table of trainable vectors, one per position. GPT-2 uses one parameter tensor of shape (1024, 768). At every forward the model looks up positions 0 through T-1 and adds the result to the token embedding. It is the simplest of the position schemes (RoPE, ALiBi, and T5 relative bias are the alternatives), and it is what the 124M reference uses.

Generation: temperature, top-k, multinomial

Generation is autoregressive: the model produces one token, appends it to the sequence, and runs again. At each step it returns logits over the full vocabulary at every position. You keep the last position only, divide by temperature, optionally mask all but the top k logits to negative infinity, softmax to get probabilities, and sample one token from that distribution.

Mermaid diagram

The three knobs change behavior in different ways. Temperature near zero collapses to greedy, and temperature one matches the model's natural distribution. Top-k of one is greedy, while top-k of forty filters the long tail. Combinations matter, and the next lesson on training uses generation as a qualitative eval signal.

Check: The Concept0/15

1. What is a learned position embedding in the 124M GPT?Define

2. What does autoregressive generation mean for a GPT model?Define

3. What is the job of the language model head?Define

4. What does weight tying mean in this model?Define

5. During generation, what does the model see once the running sequence is longer than the context length?Define

6. What does top-k filtering do before sampling?Define

7. What does temperature do at sampling time?Define

8. Why does the 124M reference use a learned position embedding?Explain

9. Why can the LM head reuse the token embedding tensor with no reshaping work beyond a transpose?Explain

10. Why does a temperature below 1 make sampling more deterministic?Explain

11. Why does generation keep only the last position's logits at each step?Explain

12. top_k_filter receives the logits [[1.0, 5.0, 2.0, 4.0, 3.0]] with top_k=2. What comes out?Apply

13. What does generate() do when called with temperature=0.0?Apply

14. Someone sets top_k=1 and temperature 1.0, then gets the same continuation from five different seeds. Why?Analyze

15. You want samples that stay varied but avoid rare junk tokens. Which sampling setting fits best?Design

Build It

code/main.py implements:

class GPTConfig, a dataclass with the 124M defaults: vocab_size=50257, context_length=1024, d_model=768, num_heads=12, num_layers=12, mlp_expansion=4, dropout=0.1, use_bias=True, weight_tying=True.
class GPTModel with token embedding, position embedding, embedding dropout, twelve TransformerBlocks, a final LayerNorm, and an lm_head that ties to the token embedding when the flag is set.
A count_parameters helper that returns the unique parameter count, so weight tying is honored in the count.
A generate function that does temperature, top-k, multinomial sampling, and a sliding window context.
A demo that builds the model, prints the parameter count next to the reference 124M, and generates a short sequence from a fixed five-token prompt ([1, 2, 3, 4, 5]) on a tiny config to show the pipeline end to end.

Run it:

bash

python3 code/main.py

Output: the parameter count alongside the 124M reference, a confirmation that the LM head and token embedding share storage when tying is on, an untied recount whose delta must equal vocab_size * d_model (asserted), the logits shape from one forward pass, generated token ids from the fixed prompt, and a sliding-window check that an 80-token random prompt on a 64-token context still generates 4 more tokens (84 total).

To keep the demo fast, the script also runs a tiny config (vocab_size=512, context_length=64, d_model=64, num_heads=4, num_layers=2, dropout=0.0) end to end and prints the generated token sequence inline. The 124M config is built (twice, once tied and once untied), but only its parameter count and one forward pass on a random 32-token input are exercised.

Check: Build It0/8

1. Why does count_parameters collect each tensor by its identity before summing?Explain

2. The demo rebuilds the reference model untied and recounts. By how much does the unique parameter count rise?Apply

3. The demo runs the 124M model on one random 32-token input with batch size 1. What shape are the logits?Apply

4. The tiny-model demo generates 12 new tokens from the prompt [1, 2, 3, 4, 5]. How long is the returned sequence?Apply

5. A 64-token-context model gets an 80-token prompt and generates 4 tokens. What happens?Apply

6. Two tiny test models use vocab_size=128 and d_model=64, one tied and one untied. What is the difference in unique parameter count?Apply

7. The same tiny model and prompt are run twice with temperature 0.9, top_k 20 and seed=123. What do the two outputs look like?Apply

8. A generation loop feeds the whole running sequence to the model and crashes with a ValueError once it passes 1024 tokens. What is missing?Analyze

Stack

torch for the tensor math, autograd, and module plumbing.
code/main.py reimplements the block pattern from lesson 34 locally.

Production patterns in the wild

Three patterns separate a model that runs from a model that ships.

Initialize the residual projections small. The attention output projection and the second MLP linear both feed straight into a residual add. Give them the same standard deviation as every other linear and the residual stream grows with depth, pushing the final LayerNorm into a hot regime. Scale their std by 1 / sqrt(2 * num_layers) and the stream stays in a sane range through twelve layers.

Cache the position id tensor. torch.arange(T) allocates fresh memory at every forward. Allocate once in __init__ for the maximum context, slice the first T entries per call, and skip the allocator round trip.

Tie weights at the parameter level. Setting lm_head.weight = token_embedding.weight shares the tensor, and copying does not. The optimizer should update one parameter and autograd should accumulate into one. Copy instead and the head drifts away from the embedding, so tying buys you nothing.

Check: Production patterns in the wild0/4

1. Why scale the attention output projection and the second MLP linear by 1/sqrt(2 * num_layers) at init?Explain

2. Why does assigning lm_head.weight = token_embedding.weight tie the weights while copying the values does not?Explain

3. In the tiny config with num_layers=2, by what factor does the init code scale the attention output projection and the second MLP weights?Apply

4. You must write a test showing a new GPT variant ties its head at the parameter level rather than copying it. Which check is strongest?Design

Use It

The model class here has the same shape as the one the next lesson trains.
Swap the learned position embedding for RoPE and you get the LLaMA family, with no change to the block or the head.
Swap GELU for SiLU and LayerNorm for RMSNorm and you get the rest of the LLaMA changes.
The generation function works with any logits source, not only this model. In lesson 37 you can pull logits from a pretrained GPT-2 file and reuse the same loop.

Check: Use It0/2

1. Which swaps turn this GPT skeleton into the LLaMA family?Explain

2. You plug a different model, such as pretrained GPT-2, into this sampling loop. What must it provide for generate() to work unchanged?Design

Exercises

Untie the LM head from the token embedding and recount parameters. Verify the delta is 50257 times 768 = 38 million.
Replace the learned position embedding with a sinusoidal table computed at construction time. Confirm the model still forwards and the parameter count drops by 786,432.
Add a greedy=True flag to generation that skips sampling and picks argmax. Confirm the sequence is deterministic across runs.
Add a repetition_penalty knob that divides the logit of any token in the prompt or generated history by a constant before softmax. Show on a fixed prompt that values above one reduce repeat counts in the output.
Add top_p (nucleus) sampling next to top_k. Two-line check that the sum of probabilities of the kept tokens exceeds top_p.

Key Terms

Term	What people say	What it actually means
Weight tying	"Tied embeddings"	The LM head and the token embedding share the same parameter tensor; saves vocab times d_model parameters and matches the GPT-2 reference
Position embedding	"Learned positions"	A separate table of shape (context length, d_model) added to token vectors; learned end to end
Sliding window context	"Context cap"	When the prompt plus generated tokens exceed the context length, drop the oldest tokens so the active window fits
Top-k sampling	"K truncation"	Keep the K logits with the highest values, mask the rest to negative infinity, softmax over the remainder
Temperature	"Sampling temperature"	Divide logits by T before softmax; T less than 1 sharpens, T equal to 1 keeps the natural distribution, T greater than 1 flattens

Author's Note

My bet is that weight tying is the first thing you should reach for at this scale and the last thing you should skip to "keep the code simple." The 38 million parameters it saves are a real fraction of a 124 million model, and the token embedding is already the matrix the head needs. I'd also build the untied version once, as in exercise 1, and check that the delta is exactly vocab times d_model. That check tells you the tie is wired at the parameter level and not copied. If the count misses the 124M reference, I'd suspect the tie or the block wiring before I'd suspect the arithmetic.