Pre-training: the full training loop

We now have all the pieces needed to pre-train our model. The complete program reads and tokenises text, constructs shifted batches, initialises a model, repeatedly optimises its weights, and saves a checkpoint that can be loaded later.

For example, this command runs two passes through the sample text, using batches of four 16-token sequences:

python code/llm/pretrain.py \
    --input code/llm/sample.txt \
    --epochs 2 \
    --batch-size 4 \
    --sequence-length 16

The model configuration and the data loader are then built from those arguments:

torch.manual_seed(1)

model_config = {
    "vocab_size": vocabulary_size(),
    "max_sequence_length": args.sequence_length,
    "model_dim": 32,
    "head_dim": 8,
    "num_heads": 4,
    "hidden_dim": 128,
    "num_layers": 2,
}
model = LanguageModel(**model_config)
batches = create_data_loader(
    args.input,
    args.sequence_length,
    args.batch_size,
)

torch.manual_seed(1) makes the random initial weights and shuffled batches reproducible.

The model dimensions and training text are tiny compared with a practical LLM, but the structure of the training loop is the same.

Run the training loop

We then pass the model, batches, and training settings into the train function:

print(
    f"Training on {len(batches.dataset)} sequences "
    f"in {len(batches)} batches."
)
train(
    model,
    batches,
    args.epochs,
    args.learning_rate,
    args.device,
)

Inside train, each epoch visits every batch once. Every batch produces one optimiser update. The next epoch asks the now-updated model to make predictions over the dataset again, so that the loss can continue to fall.

The script uses a fixed epoch count because it is a demonstration for learning purposes. Larger runs commonly add validation, learning-rate schedules, periodic checkpoints, and more sophisticated stopping criteria.

Save a reusable checkpoint

In order to reconstruct the model later, we save checkpoint files containing model.state_dict(), which holds all of the model’s weights at that point in time. We also store the dimensions used to create the model:

args.output.parent.mkdir(parents=True, exist_ok=True)
torch.save(
    {
        "model_config": model_config,
        "model_state": model.state_dict(),
    },
    args.output,
)
print(f"Saved checkpoint to {args.output}")

Generate with the trained model

By default the generate.py script loads our GPT-2 open weight file:

# implicitly loads weights/gpt2-small.pth
python code/llm/generate.py \
    "The cat sat on the"

We can instead specify a checkpoint file that our training produced:

python code/llm/generate.py \
    --checkpoint weights/tiny-teaching-model.pth \
    "The cat sat on the"

The way this works is that our code models our prompt as a “batch” of size one which it inputs into the model:

prompt_token_ids = tokenise(prompt)
token_ids = torch.tensor(
    [prompt_token_ids],
    dtype=torch.long,
    device=device,
)

Don’t expect this to produce anything sensible! We’re using a tiny input file and a tiny model. We could in principle increase the model size and amount of training text dramatically and get a more impressive result, but we would need access to powerful hardware.