LM4-L3 — Assemble a GPT Model in PyTorch

Objective

Implement LayerNorm, GELU, FeedForward, TransformerBlock, GPTModel and a greedy generation loop so that all five self-checks pass, including the exact parameter count of GPT-2 small (163,009,536 with a separate output head, 124,412,160 when it is shared).

Why a notebook? This exercise trains or runs a real PyTorch model. PyTorch does not run in the browser Python used by the site’s other labs, so you run this one on your own computer or in a hosted notebook service. The interactive browser labs of this module cover the surrounding ideas without PyTorch.

Setup and requirements

Hardware: CPU only, about 2 GB of RAM, a few minutes. CPU only; about 2 GB of RAM. The GPT-2 sized model is created on PyTorch’s meta device, which records shapes without allocating the weights, so no large memory is needed and nothing is downloaded.

Download the notebook (.ipynb) Reference solution (try the lab first)

On your own machine (Python 3.10 or newer):

python3 -m venv llm-env
source llm-env/bin/activate        # on Windows: llm-env\Scripts\activate
pip install torch numpy jupyterlab
jupyter lab assemble-gpt-in-pytorch.ipynb

In a hosted notebook service such as Google Colab or Kaggle, upload the file and run it there; PyTorch is usually preinstalled. Matplotlib is used for optional plots when it is available; the notebook runs without it.

What you will do

1. Layer norm. Normalise over the last dimension with the biased variance and eps, then scale and shift.
2. GELU and feed-forward. Write the tanh approximation of GELU and the d → 4d → d block.
3. Transformer block. Combine attention and feed-forward with two pre-LN residual connections.
4. GPT model. Add token and position embeddings, a stack of blocks, a final norm and the output head.
5. Count. Match the parameter formula on a tiny model and the exact GPT-2 small total on the meta device.
6. Generate. Write the greedy loop and see it run on random weights, with context cropping.

Cells that you must complete are marked TODO in the notebook and have a self-check right after them (an assert with a message). Run the cells from the top; a failing check tells you what is wrong.

Expected behaviour

• Self-check 1: your layer norm gives mean 0 and variance 1 per vector and equals torch.nn.functional.layer_norm.
• Self-check 2: GELU matches F.gelu(approximate="tanh"); the feed-forward block has d·4d + 4d + 4d·d + d parameters.
• Self-check 3: a block keeps the shape, does not leak future tokens, and returns its input exactly when both sublayers output zero.
• Self-check 4: logits have shape (B, n, vocab); the tiny model matches the parameter formula; GPT-2 small has 163,009,536 parameters, 124,412,160 with weight sharing.
• Self-check 5: generation keeps the prompt, is deterministic, respects the context length and returns IDs below the vocabulary size. The IDs are meaningless because the model is untrained.

How to know you are done

• All five self-checks print “passed” when the notebook is run from the top.
• You can state where the parameters of the model are and what weight sharing saves.
• The initial loss of the untrained model is of the order of ln(vocabulary size), as the last cell shows.

The site cannot run or grade the notebook for you. The self-checks inside the notebook are your verification. The reference solution was executed with PyTorch on a CPU when this module was prepared; your own environment may differ in version. No model is trained and no text is generated that carries any meaning.

If something goes wrong

• Shape error in the block: the sublayers must return the same shape as their input, (B, n, d); the residual addition needs it.
• The parameter count is off by a small number: check the biases of the feed-forward layers, the two layer norms per block and the final norm; the output head has no bias.
• Dropout makes the checks fail: the checks use model.eval or a configuration with drop_rate = 0; if you test by hand, switch dropout off.
• The meta device raises an error: it needs PyTorch 2.0 or newer. Upgrade with pip install --upgrade torch.

Theory for this lab

• LM4-T4: The Transformer Block and Assembling GPT
• LM4-T5: Output Logits, Parameter Counts and Shapes

Reference material

• Build a Large Language Model (From Scratch), Chapter 4: Implementing a GPT Model from Scratch

Build a GPT Model from Scratch labs

1. LM4-L1: Count GPT Parameters and Tensor Shapes
2. LM4-L2: Layer Normalization and Residual Numerics
3. LM4-L3: Assemble a GPT Model in PyTorch: this lab.

Next lab: LM5-L1: Cross-Entropy and Perplexity.