LM5-L3 — Pretrain a Tiny GPT and Diagnose the Run

Objective

Implement the batch sampler, the loss, the evaluation function, the training loop and a run diagnosis so that all five self-checks pass, and explain the curves of a healthy, a diverging and a memorising run.

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 is enough for the tiny configuration (about 5 to 15 minutes); a GPU makes larger runs practical. CPU is enough; the whole notebook takes about one to two minutes. On a GPU the same code runs faster, but the model is far too small to need one. Nothing is downloaded and no dataset is required: the text is generated by the notebook.

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 pretrain-a-tiny-gpt.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. Data. Read the synthetic corpus, the character tokenizer and the train/validation split (given).
2. Batches and loss. Write get_batch with targets shifted by one, and the cross-entropy loss; check that the starting loss is close to ln(V).
3. Evaluation. Write estimate_loss with the model in evaluation mode and no gradients.
4. Training loop. Write the AdamW loop with gradient clipping and periodic validation; train for 400 steps and check the curves and the generated text.
5. Sampling. Compare temperatures 0.3, 1.0 and 1.8 on the trained model.
6. Diagnosis. Write diagnose and apply it to the healthy run, a run with a learning rate of 3.0 and a memorisation run.

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: batches have shape (batch, context) and the targets are the inputs shifted by one.
• Self-check 2: the initial loss of the untrained model is close to ln(22) ≈ 3.09.
• Self-check 3: the evaluation function returns a validation loss near ln(V) for the untrained model and leaves the model in training mode.
• Self-check 4: after 400 steps the validation loss is less than half of its initial value (it reaches about 0.2, the entropy floor of the text), the training loss keeps falling, and over 80% of the words in generated text are words of the corpus.
• Self-check 5: the healthy run is diagnosed “healthy”; the learning-rate-3.0 run “diverged” or “underfitting”; the memorisation run “overfitting”. Exact numbers vary a little with the PyTorch version.

How to know you are done

• All five self-checks print “passed” when the notebook is run from the top.
• You can say why the first loss is about ln(V) and why the final loss is not zero.
• You can read the three training and validation curves and say which fix each calls for.

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 (training the healthy run in under a minute); numbers will differ slightly in your environment. The text is synthetic and highly regular, so the results do not predict what a model of this size would do on natural language.

If something goes wrong

• The loss does not fall: check that optimizer.zero_grad() comes before backward(), that you call optimizer.step(), and that the learning rate is about 3e-3.
• The initial loss is far from 3.09: the logits must be flattened to (B·n, V) and the targets to (B·n,).
• Validation loss is much lower than training loss at step 0: that is only noise from 10 random batches; use the same generator for both sets.
• It is too slow: make sure PyTorch is using several CPU threads (torch.get_num_threads()); reduce steps to 300 if needed, the self-check thresholds still hold for the healthy run.

Theory for this lab

• LM5-T2: Backpropagation and Gradient-Based Optimization
• LM5-T3: Training Loops, Validation and Checkpointing
• LM5-T4: Overfitting, Underfitting and Data Quality

Reference material

• Build a Large Language Model (From Scratch), Chapter 5: Pretraining on Unlabeled Data
• Build a Large Language Model (From Scratch), Appendix D: Adding Bells and Whistles to the Training Loop

Pretraining and Text Generation labs

1. LM5-L1: Cross-Entropy and Perplexity
2. LM5-L2: Compare Sampling Strategies
3. LM5-L3: Pretrain a Tiny GPT and Diagnose the Run: this lab.

Next lab: LM6-L1: Classification Metrics and Failure Cases.