Pretraining a GPT from scratch in PyTorch

Intro

0:00 Hello and welcome back to episode two, where we're going to pretrain our model on TinyStories, so that it goes from random tokens to writing stories. Let's begin with a quick look at where we're going to be by the end of the episode.

0:14 After a first, ten-minute training run on my laptop, our model writes stories like this one. The story isn't imaginative, and parts of it don't make sense, such as Max, who turns up from nowhere. But it reads like English, and like a story, from a model of fourteen point six million parameters.

0:35 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and the docs get a page on pretraining today. And each episode has a page here as well, with its chapters and its code.

Last time

0:54 Now, before we get to today's work, let's quickly go over what we did last time. We looked at TinyStories and at Attention Is All You Need, and wrote our decoder-only transformer in PyTorch. Token ids go into token and position embeddings, through six blocks of causal self-attention and an MLP, each with a residual connection, and out through a final norm and the head. The head shares its matrix with the token embeddings, and the whole model has fourteen point six million parameters. We printed it the way PyTorch sees it, and as a blank slate, its loss was close to nine point two, the loss of a uniform guess over ten thousand tokens, and the ids it generated were random.

Today

1:35 Today, we're going to train it, all the way from raw text to a model that writes stories. We start with a tokenizer, which splits text into words and turns them into token ids. Then we download TinyStories, build our vocabulary of ten thousand tokens from it, and store the stories as ids. We write the training loop, with batches, the loss, an optimiser and a learning rate schedule. We run it on my laptop for ten minutes, as a first run, and watch its loss fall. Then we write stories with the model, and we measure it, so we have numbers to compare against when we fine-tune it later. All right, let's get started with what the model is going to learn.

Next-token prediction

2:14 Pretraining teaches the model a single task, which is to predict the next token. We take a window of tokens from the stories, and at each position, the model sees the tokens up to there, and gives a probability to each token in the vocabulary for what comes next. We know which token actually came next, so it's our label, and the model can't see it, as the causal mask hides it. The loss is the negative log of the probability the model gave that token, averaged over the positions. If it gave the right token a probability of one, the loss is zero, and the less it gave, the larger the loss. This is the cross-entropy, and it's what PyTorch's cross entropy computes from our logits. A model that spreads its probability evenly over ten thousand tokens gives each one a ten-thousandth, so its loss is the natural log of ten thousand, about nine point two, which is close to where our untrained model was.

A tokenizer

3:08 We start with tokenizer dot py, and its docstring says how it splits text: each word, and each run of punctuation, is a token. It imports re, Python's regular expressions, and Counter, a dictionary from collections that counts things for us. Then three tokens of our own, which the model needs on top of the words. End of text, written the way GPT-2 writes it, separates one story from the next, so the model learns where a story ends. Pad fills up the shorter sequences of a batch, which fine-tuning is going to need. And unknown stands in for any word that didn't make it into the vocabulary. We keep the three in a list, with unknown first, so their ids are going to be zero, one and two. The pretokenize pattern is a regular expression that matches either a word or a run of punctuation. A word is a run of letters, digits or underscores, and whitespace matches neither, so it's dropped. The second pattern finds our three tokens in a piece of text. It escapes each of them, as square brackets and a bar mean something in a regular expression, and joins them with a bar, which means or. We note in a comment that the group around it makes re split keep the tokens it splits on, instead of dropping them. Pretokenize returns the matches of the first pattern, in order. Build vocab counts how often each token appears in the training files, a line at a time, so that a file doesn't have to fit in memory. End of text is replaced with a space first, so it's not counted as words. Then we rank the tokens by count, most frequent first, and ties go in alphabetical order, so the ids don't change from one run to the next. We keep the tokens that appear at least twice, as a token seen once is more likely a typo than a word. And the vocabulary is our three tokens followed by the most frequent words, up to the size we ask for. It also returns the counts, so we can see how much we left out. Detokenize puts the text back together, from a list of tokens. Joining the tokens with spaces leaves a space before each punctuation mark, so a regular expression takes it out, and another the spaces around an apostrophe, as in didn't. And an opening quote goes against the word after it. Then the word tokenizer class, built from a vocabulary. Id to token is the list, so a token's id is its position in it, and token to id is the reverse, a dictionary. We keep the ids of our three tokens, and len gives the size of the vocabulary. Encode first splits the text on our three tokens, and a piece that is one of them gets its id directly. Anything else is pretokenized, and each token gets its id, or unknown's when it's not in the vocabulary. And decode turns ids back into tokens, and by default, it leaves our three out before it detokenizes what's left, as a story printed for us doesn't need its padding or its end-of-text marks. The tokenizers of larger models, such as GPT-2's, use byte-pair encoding, which splits rare words into pieces, so they don't need an unknown token. Ours is the simplest thing that works for children's stories, and we can come back to it.

The data

6:13 Next, data dot py, which downloads TinyStories, builds the vocabulary, and stores the stories as token ids. It imports os, for paths, urllib's request, which is Python's built-in way of fetching a URL, and NumPy, for arrays of numbers. From our tokenizer, it takes end of text, the class and build vocab. The data folder sits next to the code, and the vocabulary size is ten thousand. The downloads are a dictionary from file names to their addresses on Hugging Face, under Ronen Eldan's account. There are TinyStories' training and validation stories, from the second version of the dataset, which GPT-4 wrote, and TinyStories-Instruct's training and validation files, which we're going to use for fine-tuning later. The first eighty megabytes of each file are plenty for a model this size. The full training set is about two gigabytes, and we'd need many hours to make use of it. Download opens the URL, and reads up to that many bytes. It decodes them as UTF-8, and errors ignore drops the bytes of a character that the cut split in half. Then it cuts the text at the last end of text, so we don't keep half a story. It writes the text to the file, and prints its size in megabytes, along with how many stories it has, by counting the end-of-text tokens. Encode reads a whole file into memory, which is fine at eighty megabytes, and tokenizes it in one go. NumPy's array turns the list of ids into one compact block, of unsigned sixteen-bit integers. Those go up to sixty-five thousand, which is plenty for ten thousand tokens, at two bytes each. Tofile writes them as raw bytes, so that training can read them straight from disk. It prints how many tokens there are, and the share of them that came out as unknown, which tells us whether ten thousand tokens are enough. Load tokenizer reads the vocabulary back, a token per line, and builds a word tokenizer from it, which is how training is going to get one. Main creates the folder if it's not there yet, and goes through the downloads, fetching each file we don't have, so running it a second time doesn't download anything again. Then the vocabulary, which we build from two files, the training stories and the training instructions. The instructions count too, and the comment says why: once the model is trained, the tokens it knows can't change, and fine-tuning is going to use their words. Build vocab gives us the vocabulary along with the count of each distinct token it saw. We open a file for writing, in UTF-8, the encoding we decoded the downloads with, and save the vocabulary to vocab dot txt, a token per line, in the order of their ids. Then we print how many distinct tokens we kept, out of how many there were. Then we make the tokenizer, and encode the training and the validation stories, which is where most of the script's time goes. And last, a sentence as the model sees it, and back. Once upon a time, a cat said hello, in quotes, and we print it, its tokens, their ids, and the text we decode from them, to check that the round trip gives us the sentence back. Git ignores the data folder, as the script can make it again.

9:22 Let's run it, and the four files come down in about ten seconds. The training file has one hundred and two thousand stories in its eighty megabytes. The vocabulary keeps ten thousand tokens, out of twenty-two thousand two hundred distinct ones. The training stories come out as about twenty million tokens, with 0.08 percent of them unknown, so the words we left out are rare ones. And our sentence comes out as eleven tokens, where the comma is a token of its own, while the exclamation mark and the closing quote are one token, as they're a run of punctuation. Once upon a time is sixty, sixty-one, seven and forty-seven, the ids we used in summary dot py, and decoding gives the sentence back.

10:03 As for why ten thousand, word frequencies fall off quickly. The hundred most frequent tokens, such as the, and, a full stop and was, make up sixty-three percent of the text. A thousand cover ninety-two percent, and ten thousand cover all but a tenth of a percent. The rest are rare names and typos, and each would cost the model three hundred and eighty-four parameters in the embeddings, for a word it sees a handful of times.

The training loop

10:30 Before the training loop, there are three small pieces that it needs. In model dot py, pick device chooses where the model runs. It tries CUDA first, which is for NVIDIA's GPUs, then MPS, the GPU in Apple silicon, which is what I'm training on, and otherwise, the CPU. Checkpoint dot py saves a trained model to a file, and loads it back, and its docstring says it's a trained model in one file: its weights, its config and its vocabulary. It imports os, asdict, which turns a data class into a dictionary, torch, and our GPT, its config and the tokenizer. The checkpoints folder sits next to the code, like the data folder. Save checkpoint first makes sure the folder for the file exists. A checkpoint holds the model's state dict, which is all its weights by name, along with the config as a dictionary and the vocabulary. So a model file has everything it needs to be used on its own. Torch save writes it, along with any extra values we pass, such as the step. Load checkpoint does the reverse, where torch load reads the file, and map location puts its tensors straight onto the device we want. Then it builds a GPT from the saved config, loads the weights into it, and makes a tokenizer from the saved vocabulary. And in sample dot py, write turns a prompt into a story. It notes whether the model was in training mode, and switches it to evaluation, which turns dropout off. It encodes the prompt and generates up to two hundred tokens. A temperature of zero point eight and the fifty likeliest tokens are a good balance between sensible and varied. It stops at end of text, switches the model back to training if it was training before, and decodes the result.

12:11 Now pretrain dot py, which trains the model for a number of minutes, as its docstring says. It imports json, for the log, math, for the cosine, along with os, sys, time, NumPy and torch. From our own files, it takes the checkpoints folder and save checkpoint, the data folder and load tokenizer, the GPT, its config and pick device, and write, for the samples along the way. The run's files go in a folder of their own, checkpoints slash pretrain. The settings are batches of forty-eight sequences, a peak learning rate of six times ten to the minus four, and a minimum of a tenth of that. The learning rate warms up over the first hundred steps, and we evaluate once per two hundred and fifty steps. Get batch picks forty-eight random starting points in the data, and takes a window of two hundred and fifty-six tokens from each, converted to sixty-four-bit integers, which is what embeddings expect. The targets are the inputs shifted by one, so position t learns to predict token t plus one, and both go to the device. Estimate loss tells us how the model is doing. It runs without gradients, switches the model to evaluation, and for the training and the validation data, averages the loss over twenty random batches, as the loss of a single batch is noisy. The validation stories are ones the model doesn't train on, so if the two losses drift apart, the model is memorising the training stories instead of learning to write. When it's done, it switches the model back to training, which turns dropout back on. The learning rate warms up linearly over the first hundred steps, as the first gradients are large and erratic, and a full-sized step on them could throw the weights off. After that, it follows a cosine, from the peak down to the minimum. We train for a set time, not a set number of steps, so progress is the fraction of the time used, and the cosine finishes as the time runs out. In main, we seed PyTorch, pick the device, load the tokenizer, and build the model, with a config whose vocabulary is the tokenizer's size. The training and validation stories are memory-mapped with NumPy, so they stay on disk, and the operating system reads in the parts we use. We print the number of parameters, the device and how many training tokens there are. We list the model's parameters, and the optimiser is AdamW, which is Adam with weight decay. Adam keeps a running average of each parameter's gradient and of its square, and scales each step by them, so each parameter gets a step size of its own. Weight decay pulls the weights a little towards zero at each step, which keeps them from growing large. We apply it to the weight matrices, and not to the biases and the layer norm gains. The betas are how quickly those running averages forget. We use zero point nine five for the second one, where Attention Is All You Need used zero point nine eight, so it adapts faster, which suits a short run. The log keeps the loss once per fifty steps, along with each evaluation. Evaluate estimates the loss, and writes forty tokens after once upon a time, so that we can watch the stories improve as it trains. It adds a row to the log, with the step, the minutes so far, both losses and the sample, and prints them on one line, with the step padded so the columns line up. It saves a checkpoint, named after its step, and the comment over it gives the reason: a checkpoint at each evaluation, to compare the model as it learns. And it writes the whole log to log dot json, so the file is up to date even if we stop the run early. Then comes the training loop, which runs until the time is up. The budget is the minutes in seconds, and we evaluate once before training, at step zero. While the time isn't up, we set this step's learning rate on the optimiser's parameter groups. We get a batch, run the model to get the loss, clear the gradients from the step before, and backward computes the new ones. Clip grad norm scales the gradients down when their combined length is over one, and above it, we write why: a rare bad batch can't throw the weights off. Then the optimiser takes its step, and the step count goes up by one. Once per fifty steps, we measure the speed in tokens a second, and log and print the loss and the learning rate. And once per two hundred and fifty steps, we evaluate, as we did at the start. When the time is up, we evaluate once more, unless the last step was an evaluation step already. Then we save the final model as model dot pt, and print how many steps it took, how many tokens it trained on, and how many minutes. The script's argument is the number of minutes, ten by default, and it can be a fraction, such as half a minute for a quick test. And git ignores the checkpoints, as they're a few hundred megabytes, and the script can make them again.

16:50 Plotted from the run's log, the learning rate climbs to its peak over the first hundred steps, and then follows the cosine down. Big steps after the warmup get the model out of its random start quickly, and small steps at the end let the weights settle into a low point of the loss, instead of bouncing around it.

A first training run

17:08 So, let's train it for ten minutes, and on my machine, through MPS, it trains at about twenty-seven thousand tokens a second. In ten minutes, it gets through twelve hundred and fifty steps and about fifteen million tokens, which is less than a pass over our training stories. The output scrolls by quickly, with the loss once per fifty steps, and at each evaluation, both losses and a sample. At step zero, the loss is about nine point three, and by step two hundred and fifty, it's already down to about three point two, and by the end, the validation loss is two point two five. The training and validation losses stay close throughout, so the model isn't memorising the training stories. Ten minutes is a starting point, as the script takes any number of minutes, so you can give it as long as your machine allows.

17:52 The loss curve, from the run's log, shows how quickly the model learns. It starts at the uniform guess, and most of the drop happens in the first two hundred and fifty steps, as the model learns which words are frequent, and then how sentences are put together. After that, each step lowers the loss by less, which is typical of training. The training batches are noisy, while the validation loss, averaged over twenty batches, is smooth, and it keeps going down until the time runs out, so more time would still help.

18:21 The log kept the sample from each evaluation, so we can read them one after another. Jq is a command-line tool for JSON, and here it prints each evaluation's step and its sample. At step zero, it's random words, as we saw. At two hundred and fifty, the first sentence is right, and then it drifts, with a bird named Tom who was very angry anymore. At five hundred, there's a boy named Tim, an idea and a line of dialogue. By seven hundred and fifty, a dog named Spot goes to the park, and by a thousand, there's a big tree with many leaves, where each sentence follows from the one before it.

Writing stories

18:58 To write stories from the command line, sample dot py gets a main, and the docstring shows how to call it, with a checkpoint and a prompt. It imports argparse, which is Python's library for command-line arguments, textwrap, for wrapping long lines, and our load checkpoint and pick device. In main, the argument parser takes its description from that docstring, and the raw description formatter keeps its layout, for when we ask for help with dash h. Each add argument call declares an argument, with its type and its default. The checkpoint is required, while the prompt is optional, and defaults to once upon a time. Dash n is how many stories to write, two by default, and tokens is how long each can be, two hundred. Then the temperature, top k and a seed, with the defaults we've been using, and parse args reads them all from the command line into a single object. We seed PyTorch with the seed, so a command writes identical stories when we run it again, which makes it easy to compare two checkpoints on one prompt. We load the checkpoint onto the best device we have, and for each story, we write it, and print it wrapped at ninety characters, with a blank line after it.

20:11 The trained model writes two stories with its default settings, and in the first, a little girl named Sue loves her toy train, asks her mom about a toy car, and they live happily ever after. In the second, a boy named Tim goes for a walk with his mom, and finds a big hill. The grammar is mostly right, and the stories have a beginning and an end. But the details drift, as the toy train becomes a toy car, and Tim wants to open the hill. The checkpoint from step two hundred and fifty is a lot less coherent, as the girl is named Amy, then Jane, and Kitty, Sue, Sam and Spot all turn up within a few lines. So the later steps are where it learned to keep track of who's who. And with a temperature of one point six, the distribution is flattened, so unlikely tokens get picked, and the sentences fall apart, although the words are still story words.

Measuring it

21:01 Samples tell us a lot, but we also want numbers, to compare models as we keep training them. Evaluate dot py measures checkpoints side by side on held-out stories, and its docstring shows how to call it. It imports argparse, json, math and os, NumPy, torch and its functional module, and our load checkpoint, the data folder and pick device. The prompt is once upon a time, there was a little. Story metrics runs a model over two hundred windows of the validation stories, twenty-five at a time. It runs without gradients, and switches the model to evaluation, so dropout is off. T is the block size, the length of a window, and we keep running totals of the loss, of the right guesses and of the tokens we've seen. The windows are back to back, from the start of the validation stories, so that each model is measured on identical text. For each batch, x and y are built like in get batch, but at fixed places, two hundred and fifty-six tokens apart. We run the model, which gives the logits and the loss. The loss is a mean over the batch's tokens, so we multiply it back by the number of tokens before adding it up. Arg max over the vocabulary gives the model's most likely token at each position, and comparing it with y counts how often that's the right one. It returns the loss, divided by the number of tokens we've seen, the perplexity, which is e to the loss, and the accuracy. Perplexity reads as how many tokens the model is choosing between, on average, so an untrained model is close to ten thousand. Next tokens shows what a model expects after a prompt. We encode the prompt, run it through the model, and take the logits at the last position, which are the scores for the next token. Softmax turns them into probabilities, top k gives the five largest along with their ids, and we return them as a dictionary from each token to its probability. Name gives each checkpoint a short name for the table, which is the run's name for its final model, and the step for the others. Main takes one or more checkpoints, which is what n args plus means, and out, an optional folder to write the results to. It picks the device and memory-maps the validation stories, like training does. For each checkpoint, it seeds PyTorch, so that each model gets identical randomness, loads the model, and keeps its metrics and its next tokens under its name. Then it prints a table, with a row per metric and a column per model, twelve characters wide, with the accuracy as a percentage and the others with two decimals. For each model, it prints its five likeliest next tokens, with a bar of hash marks for each probability. And with out, it writes the results as JSON, which is what our charts are going to read: models dot json, with a row per model and its metrics, and a file per model with its next tokens and their probabilities, each rounded to four decimals, which is plenty for a chart. Last, main runs when we call evaluate dot py from the command line. And git ignores the results folder, along with the checkpoints.

23:58 We measure all six checkpoints from our run, from step zero to the final model, and the braces in the command expand to each step's file. The untrained model has a perplexity of about eleven thousand, a little over the ten thousand of a uniform guess, and it hardly ever predicts the next token correctly. By step two hundred and fifty, the perplexity is down to about twenty-five, and the final model is at about nine point five. So at each position, it's choosing between about ten tokens, out of ten thousand, and its likeliest token is the right one about fifty percent of the time. After once upon a time, there was a little, the untrained model's guesses are random words with no probability to speak of, while the trained ones put most of it on girl and boy.

24:43 Charted over the run, the accuracy climbs quickly at first, to about thirty-six percent by step two hundred and fifty, and then slowly, to fifty percent at the end. The perplexity shows it the other way round, from about twenty-five at step two hundred and fifty down to nine point five. Both are still improving when the time runs out, so the model would keep getting better with more training, if more slowly.

25:07 After once upon a time, there was a little, this chart follows the probability each checkpoint puts on girl and on boy, and by step two hundred and fifty, the model has learned that they're the likely answers. After that, the share on girl moves around from checkpoint to checkpoint, likely as recent batches pull it one way or another. And the final model ends at about forty-one percent on girl and twenty-nine on boy, with bird, dog and cat next, at a few percent each.

25:34 Last, the docs get a page on pretraining, with the tokenizer, the data, the objective, the settings of the training loop and why we picked them, sampling and evaluation, along with what to expect from a first run. And the README now lists the commands to download the data, train, sample and evaluate.

Wrap-up

25:52 So, today we wrote a tokenizer, turned TinyStories into token ids, and trained our model for ten minutes, with a learning rate that warms up and then follows a cosine. Its loss went from a uniform guess to about two point three, which we measured as a perplexity of about ten, with its likeliest token right about fifty percent of the time. It writes stories, but as a pretrained model, it's an autocomplete. Give it an instruction, and it carries on as if the instruction were part of a story, as stories are all it's been shown.

26:23 If you train it on your own machine, I'd love to hear how far you get, and what your model writes. Let me know in the comments, and if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.