Train Your Own Small GPT

The walkthrough showed how a prediction is computed. Random weights can run through every operation correctly and still produce poor text. This lab connects that forward pass to a loss, a backward pass, and repeated parameter updates.

Here, from scratch means the small GPT starts with random weights and its architecture is implemented in the lab code. PyTorch calculates derivatives for us. The bundled text is a small teaching dataset: completing this lab demonstrates a working training pipeline, not a conversational assistant.

1. Get a Training Run Working

Download the mini GPT lab, unzip it, and open a terminal in the directory containing implementation. If you already have this repository, start at the repository root. Install uv if needed; uv sync sets up the Python environment described by the project.

cd implementation
uv sync
uv run train.py --steps 300 --device cpu --out-dir runs/tiny-gpt

The command trains on the bundled text and writes a checkpoint under runs/tiny-gpt. --steps 300 asks for 300 optimizer updates. It is a starting experiment, not a promise that training will converge after that many updates.

Before editing the code, run the four phases below and match each phase to the script:

What to look for: The forward pass computes probabilities, the loss scores the targets, backward computes gradients, and the optimizer changes parameters. The next forward pass uses those changed parameters. The animation illustrates the pattern; your terminal reports the measured values for your run.

First checkpoint: Find the initial validation loss and a later validation loss in the run output. Keep both values even if the later value is worse. That is a result to investigate.

The output directory also contains metrics.csv with step, training loss, and validation loss; loss.svg plots those measurements. samples.json records text from before and after training, and config.json records the run settings. Keep these files with the checkpoint.

A reference CPU run with the bundled corpus, default settings, seed 42, Python 3.12, and PyTorch 2.14 recorded:

These measurements show learning on this small corpus. Your values can differ with library versions and devices. The generated text still contains broken words; improving next-character prediction on a few short stories does not produce a general chat model.

2. Check What the Model Is Asked to Predict

Suppose the text is The cat. A four-character training window can use:

input:   T h e [space]
target:  h e [space] c

Each position predicts the character immediately after it. The first target is h, not T. The lab uses characters so you can inspect the whole vocabulary; the same shift applies to the subword tokens from Tokenization.

Think of the input and target as two copies of a sentence, with one copy moved one place left. This describes the alignment, but the model still sees only its allowed prefix at each position.

The code splits the source text into training and validation portions before constructing overlapping windows. Otherwise, nearly identical windows could appear in both splits. The tokenizer learns its vocabulary from the training portion; an unseen validation character maps to the unknown ID.

import torch
from data import CharacterTokenizer, get_batch

text = "The cat sat on the mat. "
tokenizer = CharacterTokenizer.fit(text)
stream = torch.tensor(tokenizer.encode(text), dtype=torch.long)
x, y = get_batch(stream, batch_size=2, block_size=4)
print(x.shape, y.shape)  # both: (2, 4)
print(tokenizer.decode(x[0].tolist()))
print(tokenizer.decode(y[0].tolist()))
assert torch.equal(x[:, 1:], y[:, :-1])

3. Follow One Update Through the Code

For a batch of 2 sequences, 4 positions each, and a vocabulary of 30 characters, the model receives IDs shaped (2, 4) and produces scores shaped (2, 4, 30). There are 8 predictions, each choosing among 30 characters.

These raw scores are called logits. The loss function handles the probability calculation internally. Pass logits directly to cross-entropy; do not apply softmax first.

optimizer.zero_grad(set_to_none=True)
logits, _ = model(x)                       # shape: (batch, positions, vocabulary)
loss = torch.nn.functional.cross_entropy(
    logits.reshape(-1, logits.size(-1)),   # shape: (batch * positions, vocabulary)
    y.reshape(-1),                        # shape: (batch * positions,)
)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

If these roles feel unfamiliar, revisit Neural Network Basics and the Training Loop. You do not need to derive every attention gradient before running this experiment.

4. Make a Small Test Pass Before a Larger Run

A model that cannot learn one fixed batch is unlikely to learn an entire corpus. Use the diagnostic mode to repeatedly update on the same batch:

uv run train.py --overfit-one-batch --steps 200 --device cpu --out-dir runs/one-batch

The fixed batch loss should fall substantially with enough updates. This deliberately tests memorization; its validation result is not evidence of generalization. If it fails, inspect the target shift, gradients, and learning rate before adding layers or more data.

Next, train the bigram baseline, which predicts the next character using only the current character:

uv run train.py --model bigram --lr 0.03 --steps 300 --device cpu --out-dir runs/bigram

Compare it with the small GPT using the same source text and split. The baseline command uses a larger learning rate because its small lookup table trains differently from the transformer; record both settings in your report. A longer-context model has more information available, but it still needs enough useful data and successful optimization to beat the baseline.

5. Generate, Stop, and Resume

Load the trained weights in a separate process:

uv run generate.py --checkpoint runs/tiny-gpt/checkpoint.pt --prompt 'The ' --max-new-tokens 100

Read the output alongside the validation loss. Recognizable words can come from memorized training passages; a low loss on a tiny held-out split does not establish broad language ability.

checkpoint.pt contains the latest state. best.pt contains the state with the lowest measured validation loss; pass that path to generate.py to compare an earlier, better validation result with the final one.

Continue the same training run for 100 additional updates:

uv run train.py --resume runs/tiny-gpt/checkpoint.pt --steps 100

A checkpoint stores the model weights, optimizer state, step count, tokenizer, and configuration needed to continue the experiment. The optimizer has its own memory of earlier gradients. Loading weights alone is sufficient for generation, but does not fully resume a training run.

Second checkpoint: Confirm the resumed step count continues from the saved count. Then load the new checkpoint with generate.py again. Keep the data and configuration unchanged when testing continuation.

6. Submit an Experiment Someone Else Can Check

Save these items together:

1. The training command, seed, data source, and split method.
2. Initial and final training/validation measurements for the baseline and GPT.
3. The one-batch test result and one problem you checked while debugging.
4. Generated text from fixed prompts, with the checkpoint and sampling settings recorded.
5. The checkpoint and commands for loading it and resuming training.

Write a short explanation of whether the GPT improved on the baseline, and what the validation evidence supports. If it did not improve, describe one change to test next while holding the other settings fixed. This report is the completion task for the core course.

Quick Summary

You now have a training experiment to inspect. Next, compare how the same transformer components support different objectives in AI Models →. After GPT, you can reuse this training workflow in the optional open-model fine-tuning lab.