SFT: fine-tuning a GPT to follow instructions

Intro

0:00 Hello and welcome back to episode three, where we're going to teach our model to follow instructions, with supervised fine-tuning. Let's begin with a quick look at where we're going to be by the end of the episode.

0:13 Here, we ask our fine-tuned model for a story with a dialogue, and the words cake, share and jump. It writes about a boy named Timmy, with a cake, a lot of jumping, and his mom saying it's okay to share, over a few lines of dialogue. It wanders, as Timmy finds a big puddle and jumps up and down, and they end up playing together in a big pond. But the story uses all three words, and it has the dialogue we asked for.

0:39 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 fine-tuning today, and one on watching training. And each episode has a page here as well, with its chapters and its code.

Last time

0:59 Now, before we get to today's work, let's quickly go over what we did last time. We wrote a word-level tokenizer, with a vocabulary of the ten thousand most frequent tokens in TinyStories, and stored the stories as token ids. Then we pretrained our model for ten minutes, to predict the next token, with AdamW, a warmup and a cosine learning rate. It went from random words to short stories that mostly hold together, and we measured it on held-out stories, with a validation loss of about two point two five and a perplexity of about ten. But it's an autocomplete, and today we're going to teach it to follow instructions.

Where SFT fits

1:37 This is the pipeline behind assistants such as ChatGPT, in a much simpler form. Pretraining on a lot of text gives a model that knows the language, and continues whatever it's given. Supervised fine-tuning then trains it further on examples of an instruction followed by the answer we want, written by people, or here, generated by GPT models. It uses pretraining's objective, next-token prediction with cross-entropy, on different data and with one change to the loss. After that, preference optimisation, such as DPO, teaches it to prefer the better of two answers, and reinforcement learning, such as PPO, trains it to maximise a reward. Today, we're going to do the supervised step, which is what teaches our model to follow instructions. All right, let's get started with the data we're going to fine-tune on.

TinyStories-Instruct

2:25 Our data is TinyStories-Instruct, a variant from the TinyStories paper, where each story comes after a set of instructions. There are four kinds of instruction: a list of words to use in the story, a sentence that should appear somewhere in it, a list of features, such as a dialogue, a bad ending or a moral value, and a short summary. The words and features are the ones each story was generated from. GPT-3.5 wrote the summaries, and the sentence is picked at random from the story. Each story gets a random subset of the four, so the model has to cope with any combination.

2:58 In the file we downloaded in episode two, the first story's instructions are features, dialogue, the words quit, oak and gloomy, and a summary, followed by story and a colon, where the story starts. The paper describes the instructions as coming before the story. The fifth document, which we print with a line of Python, doesn't follow that. It asks for the words help, mud and bossy, then comes the story, and its summary comes after it, at the end. In the training file, about thirty percent of the documents with a summary have it after the story. If we trained on the text as it is, the model would learn to write a summary after its stories, and our instructions would be missing their summary almost a third of the time. So we're going to parse each document into its fields, and put them back together in one order.

3:42 First, let's see what the pretrained model does with an instruction. We ask for a story with the words dog, ball and happy, and it carries on as if it were in the middle of a story, starting on a line of dialogue, with Mia and someone called Max. It uses happy, and the second story uses ball, as attention can copy tokens from the prompt, but neither starts at the beginning of a story, and neither uses the word dog. It hasn't been shown an instruction followed by its answer before, so the instruction is just more text to continue.

A loss on the story alone

4:13 The loss is the one from pretraining, with one change. The model still sees all the tokens before each position, the instruction along with the story so far, but we score the positions in the story, S, and no others. In pretraining, the model learned to predict each token of the text. If we did that here, a lot of the loss would go into predicting the instructions, such as which words someone might ask for, which isn't what we want it to learn. So it reads the instruction, and it's scored on the story that follows it. In code, that's a target of minus one at each position we don't score: the instruction, and the padding that brings the examples of a batch to one length. The story's targets are its tokens, ending with end of text, so the model also learns where a story should stop.

Instructions and stories

5:00 We start sft dot py, and its docstring says what it does: the pretrained model learns to write the story an instruction asks for, and we run it with a number of minutes. It imports re, for regular expressions, and end of text from our tokenizer. Fields lists the parts an instruction can have, in the order our prompts give them, before the story. The field pattern finds a field's name at the start of a line, followed by a colon. Multiline makes the caret match at the start of each line, and not at the start of the text alone. Parse splits a document on those names, wherever they are in it, using the pattern's split. With the group in the pattern, the split alternates names and texts. So pairing the odd items with the even ones after them gives a dictionary from each name to its text, with the whitespace stripped. Whatever order the fields came in, we now have them by name. Instruction writes them back out in our order, a line per field the document has, and ends with story and a colon. The model writes its story from that last line on. And as each prompt gives the fields in one order, the prompt is easier to learn. Build examples reads a file, and splits it into documents on end of text. It drops the empty ones, such as the one after the last end of text, and keeps the first limit of them, if we give one. For each document, we parse it, and skip it if it has no story. The prompt is the instruction's token ids, and the story is the story's ids followed by end of text, so the model learns to stop when the story's done. The example is the two together, cut to the context window plus one, as the inputs and the targets are shifted by one. We keep an example when some of its story made it in, along with the length of the prompt, which is where the story starts. An instruction that fills the window leaves no story to learn from.

Padding and masking

6:51 First, there's a change to model dot py, in the loss. Cross entropy gets ignore index minus one, so a target of minus one is skipped, and doesn't count towards the loss or the gradients. As the comment says, a target of minus one marks a position with no token to learn, such as padding. In sft dot py, we import NumPy and torch, and a batch is forty-eight examples, like in pretraining. Make batch turns a list of examples into tensors. Examples have different lengths, and a tensor has to be a rectangle, so we find the longest. X is filled with the pad id, and y with minus one, at that length minus one. For each example, x gets its ids, without the last one, and the rest of its row stays padding. Y gets the story's ids, from the position before the story starts. So the story is what's learned: the last position of the instruction predicts the story's first token, while the instruction and the padding keep their minus one, which the loss skips. The instruction is still there in x, so attention sees it, and that's how the story can follow it, while the loss doesn't score the model on predicting the instruction. Random batch picks forty-eight examples at random, and makes a batch of them. And estimate loss averages the loss over twenty random batches, without gradients and in evaluation mode, like in pretraining, but over our examples this time. It returns the average, which is the number we're going to watch go down as the model fine-tunes.

The fine-tuning loop

8:18 First, write, in sample dot py, gets an argument, with prompt, which decides whether the text it returns starts with the prompt. It's true by default, so our other calls work as before. For fine-tuning, we want the story on its own, so with prompt false returns the tokens after the prompt. Back in sft dot py, we import what the training loop needs: json, os, sys, textwrap and time, our checkpoints, the data folder, pick device and write. The run goes in checkpoints slash sft, and the learning rate is two times ten to the minus four, a third of pretraining's peak, and we evaluate once per two hundred and fifty steps. The demo is the instruction we write a story for at each evaluation: a dialogue, with the words dog, happy and park, and a summary, a dog makes a new friend at the park. In main, we seed PyTorch, pick the device, and load the pretrained model from its checkpoint, along with its tokenizer, as fine-tuning has to use the vocabulary the model was trained with. The block size comes from the model's config, as the examples are cut to it. We build the training examples from the instructions' training file, and the held-out ones from the first five thousand documents of the validation file, which is plenty for a loss estimate. Building them takes a few seconds, as each document is parsed and tokenized, and we print how many there are of each. We list the model's parameters for the optimiser, like in pretraining. The optimiser is AdamW, with the weight decay groups from pretraining, and a lower learning rate than pretraining's. A big one would undo a lot of what pretraining taught it, which is called catastrophic forgetting. The matrices get a weight decay of zero point one, and the biases and the layer norm gains get none, as before. Evaluate measures the loss on the training and the held-out examples, and writes a story for the demo instruction, up to a hundred and twenty tokens, without the prompt. It logs them, with the step and the minutes so far, and prints the losses, with the story wrapped at eighty-six characters and indented under them, so it's easy to tell apart from the numbers. It makes sure the run's folder exists, and writes the whole log to log dot json. The budget is the minutes in seconds, we note when we started, and we evaluate once at step zero, before any fine-tuning, to have a baseline. While the time isn't up, the loop is pretraining's, without the learning rate schedule: a random batch of examples, the loss, backward, the gradients clipped at one, and the optimiser's step. Once per fifty steps, we log and print the loss, and once per two hundred and fifty, we evaluate. When the time's up, we evaluate once more if we need to, save the model in checkpoints slash sft, and print how many steps it took, and how many minutes.

A short fine-tuning run

11:01 So, let's fine-tune for ten minutes, on about seventy-seven thousand training examples. At step zero, before any fine-tuning, the loss on held-out stories, given their instructions, is two point five. That starting loss is already quite low, as these are TinyStories' stories, which our model knows how to write, and the instruction isn't part of the loss. After one thousand one hundred and sixty-seven steps, the held-out loss is two point zero two, and at each evaluation along the way, the output has the story the model wrote for the demo instruction.

11:32 Jq prints those stories from the log, one per evaluation, for a dialogue with the words dog, happy and park, about a dog that makes a new friend at the park. At step zero, the pretrained model carries on with a dialogue between Amy and Tom, with no park and no new friend. By two hundred and fifty steps, the story starts at the beginning, and there's a dog at the park making a new friend, even though the dog's name comes out as the park. From there on, the stories stay with the instruction. By seven hundred and fifty, a happy dog in a big park meets a new friend, which is what the summary asked for. The model repeats the word dog a lot, as in the dog and the dog played together. So it's following the instruction, but its writing hasn't got any better.

12:15 The loss on held-out stories, given their instructions, goes from about two point five to two point zero. As the model already writes stories like these, what it's learning here is to use the instruction: which words to use, which features, and what the summary says should happen. The training and held-out losses stay close, so it isn't memorising the training examples.

Before and after

12:35 Now we give the fine-tuned model the instruction from before, the words dog, ball and happy. Both stories start at the beginning, with once upon a time, and use the three words, with a happy dog and his ball. The writing is still simple and repetitive, as it's still a model of fourteen point six million parameters, and ten minutes of fine-tuning doesn't make it a better writer. With a bad ending, and a summary where Tom loses his red kite on a windy day, it writes about Tom and his red kite, and Tom crying because he can't see it. So it follows the summary in part, and it runs out of tokens before the story ends, which we're going to see in the numbers too.

Measuring it

13:14 Samples are a good first look, but we want numbers to compare the two models, so evaluate dot py now measures how well a model follows instructions too, as its docstring says. It imports re, and from sft dot py, build examples, instruction, make batch and parse, along with end of text. We're going to write stories for fifty held-out instructions, up to two hundred tokens each. Fifty is how many sets of instructions the TinyStories paper used for its own evaluation, where GPT-4 graded the stories. We're going to count instead, which is cruder, but free and repeatable. Load instructions parses the documents of the validation file, the ones the model hasn't trained on. It keeps the first fifty documents that have a story and ask for words, so we can check whether the story uses them. Each comes back as the instruction's text and its list of words. Instruction metrics writes a story for each of the fifty instructions. For each one, we encode the instruction, and generate up to two hundred tokens, with the sampling settings we use elsewhere, stopping at end of text. We keep the ids after the prompt, and decode them into the story. Then we count the requested words that appear in the story, with a regular expression that looks for each one as a whole word, ignoring case, with an optional ending. Above it, we write how it matches: a word counts in its plural, or with d, e d or i n g on the end, so play counts in played. We add up how many of the words appear, count the stories that end with end of text, and keep their lengths. It's a simple rule, and it misses irregular forms, such as ran for run, which is fine for a rough measure. The instruction loss is the loss on the stories of held-out examples, in batches of twenty-five, through make batch, so it's the loss fine-tuning trains on, measured on stories the model hasn't seen. It returns the share of words used, the share of stories that stopped on their own, and their average length. In main, we load the fifty instructions once, for all the models, from the validation file. Then, for each checkpoint, we seed PyTorch, load the model, and build five hundred held-out examples with its tokenizer and its block size, and merge the story metrics with the instruction metrics. The table's first column is wider, for the longer names, and words used and stopped show as percentages too. The results go to JSON as before, so the charts can read the new metrics along with the old ones.

15:42 We compare the pretrained model with the fine-tuned one, and on the fifty instructions, the fine-tuned model's stories use twenty-seven percent of the requested words, against two point seven percent before, about ten times as many. The loss on held-out instructions' stories goes from two point six zero to two point one three. Fewer of its stories end within two hundred tokens, forty-eight percent against ninety-two, and they're about twice as long, at one hundred and seventy-eight tokens against eighty-five. And on plain stories, the loss goes up a little, from two point two five to two point three eight.

16:15 The words are what fine-tuning was for, and that's the clearest change. Fewer of the fine-tuned model's stories end with end of text, and the length chart on the right shows why. The pretrained model stops more often, but its stories are short, as it carries on from the middle of a story and soon wraps it up. The fine-tuned model writes whole stories, closer in length to the instructions' stories, and those often run past two hundred tokens, so a lot of them are cut off before they end. On top of that, stories longer than our context window lost their end of text when we cut them to the window, so the model saw fewer story endings than the dataset holds.

16:53 The two losses show what fine-tuning changed, and on the instructions' stories, the loss fell from two point six zero to two point one three, as the instruction now helps the model predict its story. On plain stories, it rose a little, from two point two five to two point three eight. The fine-tuning stories come from the dataset's first version, which GPT-3.5 partly wrote, while our validation stories are from the second version, written by GPT-4 alone. So the model has likely moved a little towards the first version's stories. Fine-tuning made it better at what we trained it on, at a small cost elsewhere, and a lower learning rate or fewer steps would keep that cost smaller.

17:33 With that, the docs get a page on supervised fine-tuning, with how TinyStories-Instruct's fields become instructions, the loss on the story alone with a table of the targets, the training loop, and the instruction metrics, along with what to expect. And the README now has the commands to fine-tune, write a story for an instruction, and compare the two models.

Better logs with TensorBoard

17:52 So far, each run has written its numbers to log dot json, which we read with jq, and which our charts are plotted from. That works, but each run overwrites the log of the one before, and the numbers are just a file until we write code to look at them. So let's log to TensorBoard instead, a tool from TensorFlow for watching training, which PyTorch can write to. You'll notice the files don't start with a docstring any more, as the docs say what each file is for. In checkpoint dot py, next to the checkpoints folder, there's a runs folder, where each training run logs to a folder of its own. In pretrain dot py, json goes, and we import SummaryWriter from PyTorch's TensorBoard utilities, along with the runs folder. A writer replaces the log, and its comment notes that the run's numbers go to TensorBoard, in a folder named after the run and when it started, as strftime formats the time as the year, month and day, then the hours, minutes and seconds. Evaluate doesn't take the elapsed time any more, as TensorBoard records the time of each value on its own. Add scalar is how the writer logs a number: it takes a name, the number, and the step it belongs to. So the training loss goes under loss slash train, and the validation loss under loss slash valid, and the slash groups the two together in TensorBoard. Add text logs the sample, and flush writes it all to disk now, so that TensorBoard shows the evaluation while the run goes on, and the lines that wrote log dot json go. The first evaluation, at step zero, takes just the step. Once per fifty steps, the batch's loss, the learning rate and the speed are three more scalars. And at the end, we evaluate once more unless the last step was an evaluation step, which we can tell from the step now, and close the writer, which writes whatever is left. SFT dot py gets the changes we made to pretraining. Json goes, and we import SummaryWriter and the runs folder. The writer logs to a folder named after sft and the time it started. Evaluate takes just the step, logs the training and held-out losses under the names pretraining uses, and the sample as text, and flushes them, and the lines that wrote log dot json go. The first evaluation takes the step, once per fifty steps the batch's loss is logged, and at the end, we evaluate once more if we need to, and close the writer, as pretraining now does too. TensorBoard joins the requirements, as SummaryWriter needs it installed, and it's also the app we open to watch the runs, and git ignores the runs folder, like the checkpoints.

20:24 Let's fine-tune again, for three minutes this time, so we have a run to look at. It starts from our pretrained model, with the seed from before, so step zero matches the first fine-tuning run, and the output looks as it did, as the prints haven't changed. The numbers now go to the runs folder too, where the writer adds them as the run goes. In three minutes, it gets through about three hundred and fifty steps, which is plenty to see the loss come down.

20:52 TensorBoard runs as a small web server: tensorboard dash dash logdir runs, and we open it in the browser, on port six thousand and six. The time series tab has a chart for each scalar we logged, grouped by the part of the name before the slash, so the three losses sit together. The loss of a training batch jumps around, as each batch has different stories, so TensorBoard draws a smoothed line over it, with the raw values faint behind. The held-out loss, from the evaluations, goes down steadily, as it did in our first run. The smoothing slider sets how much it smooths, and hovering over a chart shows the values at each step. And the text tab has the samples, one per evaluation, so we can read how the stories changed without leaving the browser. With several runs in the runs folder, they'd all show up here, each in a colour of its own, which is how we'd compare two learning rates, say.

21:47 TensorBoard is good for watching a run, but sometimes we want the raw numbers, to compare runs our own way, or to make a chart for a write-up. The event files can be read in Python, so plot dot py reads them, and plots them with matplotlib, Python's most used plotting library. It imports argparse and os, matplotlib's pyplot, and the event accumulator, from TensorBoard's own code. The colours are the series': periwinkle for the training batches, amber for validation, on the near-black slate, with a dim grey for the frame of each chart. Read run makes an accumulator for a run's folder. The accumulator reads a run's event files, the ones TensorBoard reads too, and reload loads them. Tags lists the names that were logged, by kind. For each scalar, we get a list of events, each with the wall time, the step and the value. We keep two lists, the steps and the values, which is what matplotlib wants, one list for each axis. Then come the samples, which we logged with add text, and the comment over them explains how they're kept: TensorBoard stores text as a tensor of strings, one per evaluation, under sample slash text summary, so for each evaluation, we take its step and decode its string from bytes. Plot starts with matplotlib's dark style, which suits our dark background. The learning rate gets a chart of its own, as it's on a scale of its own, thousands of times smaller than the loss. It's pretraining that logs it, since fine-tuning keeps it fixed. Subplots makes a figure with one or two charts side by side, six inches wide each, and squeeze false keeps the charts in a grid even when there's just one. On the loss chart, plot draws the training batches as a line, where the star unpacks our steps and values into x and y, and the validation loss as a line with a dot at each evaluation, and each gets a label, for the legend. Then the title and the x axis's label, and the learning rate's chart, if there is one, drawn as a single line, with its own title. Each chart gets the slate background, a dim frame and a faint grid, the figure is titled with the run's name, and savefig writes it as a PNG, at a hundred and fifty dots per inch. Main makes an argument parser, with a one-line description of the script. It takes one or more run folders, and a folder for the charts, results by default, and it makes sure that folder exists. Then we go through the runs, one at a time, and for each, the name is the last part of its folder. We read it, and print each scalar's name, how many points it has, and its last value and step, then each sample's step and its first eighty characters, and plot it to a PNG named after the run, in the charts folder, so that a single command charts several runs at once. And matplotlib joins the requirements, next to TensorBoard and torch.

24:36 Let's run it on the three-minute run's folder, and it finds three scalars: the training and held-out losses, with three points each, at step zero, two hundred and fifty and the last step, and the batch loss, with six points, once per fifty steps. Then come the three samples, cut to eighty characters, and it writes the chart to the results folder, named after the run.

24:57 The chart is in the series' colours, and the held-out loss comes down steadily over the three evaluations, while the loss of a single batch jumps around above it, as it's one batch of forty-eight examples, not an average over twenty batches. A pretraining run would get the learning rate's chart next to it. From here, we can make any chart we like from our runs, such as the held-out losses of two runs on one chart, to compare them.

25:21 And the docs get a page on watching training. It has a table of the names each run logs its numbers under, and how often, the command to start TensorBoard, and a few lines of Python that read a run's held-out loss back, along with how to plot the runs with plot dot py. The home page and the pretraining page link to it. The README has the commands to start TensorBoard and to plot the runs.

Wrap-up

25:45 So, today we fine-tuned our model to follow instructions. We parsed TinyStories-Instruct's fields into one order, so each summary lands in the instruction, and masked the loss, so the model is scored on the story and not on the instruction. In ten minutes, at a third of pretraining's learning rate, the share of requested words its stories use went from about two point seven to about twenty-seven percent, at a small cost on plain stories. And our runs now log to TensorBoard, where we can watch the losses while a model trains, and read them back in Python to plot them. It's still a small model, and its stories still wander, but it now writes the story an instruction asks for, instead of continuing the instruction as text.

26:29 Let me know in the comments what you'd like the model to learn next, and which parts of the pipeline you'd like us to go deeper on. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.