LoRA: fine-tuning a GPT with 2% of its weights

Intro

0:00 Hello and welcome back to episode nine, where we're going to fine-tune our model on instructions again, this time with LoRA, which trains about two percent of its weights. Let's begin with a quick look at where we're going to be by the end of the episode.

0:14 By the end of today, our model follows instructions through an update to about two point two percent of its weights, kept in a file of about one point three megabytes, next to the model's fifty-six. With the update merged into the model, and asked for a story with dog, ball and happy, it uses all three.

0:32 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 LoRA today, with what it saves and what it learns. And each episode has a page here as well, with its chapters and its code.

0:53 Now, before we get to today's work, let's quickly go over what we did last time. We trained our fine-tuned model with reinforcement learning, with PPO. It wrote stories as it trained, a reward scored the share of the requested words each one used, and PPO made the stories that scored well more likely, with a critic estimating the reward to come. A KL penalty kept the model near a frozen copy of where it started, and its weight set how far the model could move to raise the reward, further with a small weight than with a large one. Fine-tuning in episode three, DPO, and both of PPO's runs each trained all fourteen point six million of the model's weights, and each saved a full copy of the model.

Low-rank adaptation

1:35 Full fine-tuning updates each weight of the model, so AdamW keeps two numbers of its own for each of them, and each fine-tuned version is a full copy of the model. Today, we're going to fine-tune with LoRA, low-rank adaptation, which freezes the pretrained weights, and trains a small update beside some of the layers, as the product of two thin matrices, B and A. We fine-tune on the instructions from episode three, from the pretrained model, for ten minutes, like we did then, so that the two can be compared. And we measure what it saves, in weights trained, in the file we keep and in the optimiser's memory, and what it learns, in the losses and the words used. All right, let's get started with the paper that introduced it.

2:15 The paper is LoRA, low-rank adaptation of large language models, by Edward Hu and others at Microsoft, from 2021, and one of its authors, Yuanzhi Li, is also one of the two authors of TinyStories. Its abstract starts from GPT-3, with a hundred and seventy-five billion parameters, where keeping a fine-tuned copy of the whole model for each task is too expensive. LoRA freezes the pretrained weights and adds small trainable matrices beside the layers, which cuts the trainable parameters by about ten thousand times and the GPU memory it needs by three times, with no extra time when the model runs. In figure one, the pretrained weights, W, are frozen, and the input also goes down through A, to r numbers, and back up through B, and the two paths are added. Section four says why so few numbers can be enough. Earlier work found that fine-tuning a pretrained model needs far fewer dimensions than it has weights, so the authors expect the change to each weight matrix to have a low rank too.

The update, B A

3:11 A linear layer computes W x, and full fine-tuning changes each of W's numbers. LoRA leaves W as pretraining left it, and adds an update to the output, B A x, scaled by alpha over r. A takes the layer's input down to r numbers, where r is the rank, and B takes those r numbers back up to the size of the layer's output. So B times A has the shape of W, but its outputs can move along r directions and no more, which is what having a rank of r means. For one of our attention's projections, three hundred and eighty-four by three hundred and eighty-four, W holds a hundred and forty-seven thousand weights, while at a rank of eight, A and B hold about six thousand between them, about four percent of that.

3:57 B starts at zero, so the update is zero at the first step, and the adapted model starts out as the pretrained one. A starts at random, as with both of them at zero, neither would get a gradient, and the update would stay at zero. Alpha over r scales the update, here sixteen over eight, which is two. The paper sets alpha once and leaves it, so that trying another rank doesn't mean tuning the learning rate again. Once it's trained, the update folds into the weight. W plus alpha over r times B A has W's shape, so the merged model is a plain model again, as fast as the original, with no extra layers to run.

4:35 We adapt the linear layers inside each block, starting with attention's query, key, value and output projections, at about six thousand weights each. Then the MLP's up and down, at three hundred and eighty-four by fifteen hundred and thirty-six, which get about fifteen thousand each. That makes about fifty-five thousand per block, across six blocks, for about three hundred and thirty thousand weights in all. The embeddings, the layer norms and the head stay as pretraining left them, so all the change from fine-tuning goes through the adapted layers.

LoRA in PyTorch

5:06 So, lora dot py is a new file, for fine-tuning with LoRA. It imports math, for a square root further down, os, sys, textwrap and time, torch, its neural network module, and TensorBoard's writer. From our own files come the checkpoints and runs folders, with loading and saving a checkpoint, the data folder, pick device and write. And from sft dot py come the demo instruction, how often to evaluate, and the functions that build, batch and score the instruction examples, so LoRA learns from the data episode three fine-tuned on. The run goes in checkpoints slash lora, the rank is eight, and alpha is sixteen, so the update is scaled by two. The learning rate is one times ten to the minus three, five times fine-tuning's, and the comment gives the reason: the adapters start from nothing, and there are few of them. Targets names the layers we adapt: attention's query, key, value and output projections, and the MLP's up and down. LoRA linear is a module that wraps one of those layers, and its comment says what it adds: a frozen linear layer, plus a trainable low-rank update, W x plus alpha over r B A x. It keeps the layer it wraps as base, and A is a parameter of the rank by the layer's input size, on the layer's device, and kaiming uniform fills it with random values, PyTorch's default for a linear layer's weights, which the paper's own code uses for A too. B is the layer's output size by the rank, in zeros. B starts at zero, so the adapted model starts as the pretrained one. And scale is alpha over the rank, which the update is multiplied by. Forward returns the base layer's output, plus x through A and then B, times the scale, with A and B transposed, as a linear layer does with its weights. That way, the full matrix B A isn't made, as each token goes down to eight numbers and back up. Add lora freezes the whole model first, with requires grad false. Then, in each block, for the attention and the MLP, each target layer the module has is swapped for a LoRA linear around it, with setattr. The A and B made this way come after the freeze, so they're all the optimiser gets to train. Merge runs with no gradients, and folds each update into its weight, W plus alpha over r B A, and puts the plain layer back. For each LoRA linear it finds, the base weight gets the scale times B times A added in place, and setattr puts the base layer back where the wrapper was. Main seeds PyTorch, loads the pretrained model from episode two, as full fine-tuning did, and adds LoRA to it. It counts the trainable parameters, and all the parameters, which now include the adapters, and prints both, with the share. Then come the training examples, and five thousand held-out ones, as sft dot py builds them. AdamW gets the trainable parameters alone, with no weight decay, as they're few and start small, and the writer logs to a folder named after lora and the time. Evaluate estimates the loss on the training examples and on the held-out ones, each over twenty batches, and writes a story for the demo instruction, up to a hundred and twenty tokens, without the prompt, with the temperature and top k we use elsewhere. It logs both losses and the story to TensorBoard, flushes them to disk, so we can watch the run as it goes, and prints the losses with the story under them, like fine-tuning does. The budget is the minutes in seconds, and we evaluate once at step zero, where the model is still the pretrained one. Then the loop is fine-tuning's: a random batch of forty-eight examples, the loss, backward, the gradients clipped at a norm of one, and the optimiser's step. Each fifty steps, the batch's loss goes to TensorBoard and to the terminal, with the minutes so far, and each two hundred and fifty, we evaluate, until the time is up. Backward still runs through the frozen layers, as the adapters in the first blocks need the gradient from the blocks above them. After the loop, we evaluate once more if the last step wasn't an evaluation, and close the writer. At the end, we add up the size of the optimiser's state, and above it, we note what it holds: AdamW keeps two numbers per trainable parameter, its running averages of the gradient and of its square. We print it next to eight bytes for each of the model's parameters, which is what AdamW would keep in full fine-tuning. The adapters, each trainable parameter by its name, go to adapters dot pt, with the rank and alpha, and we print the file's size. Save checkpoint writes the merged model to model dot pt. Merged, it's a plain model again, which sample dot py and evaluate dot py load like any other. Last, main runs with the minutes we pass, ten by default.

Ten minutes of LoRA

9:47 Let's run it for ten minutes, and it trains about three hundred and thirty-two thousand of fourteen million nine hundred thousand parameters, about two point two percent, where the total counts the adapters too. At step zero, the held-out loss is about two point four seven, the pretrained model's, as B is zero. It falls to about two point two six by step five hundred, and then moves little, to end at about two point two four after one thousand two hundred and eighty-two steps. That's about ten percent more steps than full fine-tuning managed in its ten minutes, as the frozen weights need no gradients of their own, and AdamW updates few numbers. At the end, AdamW's state takes about two point seven megabytes, against about one hundred and nineteen for all the parameters, and the adapters' file is about one point three five megabytes.

10:33 The chart puts LoRA's held-out loss next to full fine-tuning's from episode three, step by step, from where the pretrained model starts. By step two hundred and fifty, full fine-tuning is a little lower, and after that, it keeps going down, to about two point zero three at step a thousand, while LoRA levels off at about two point two seven. So in ten minutes, LoRA's held-out loss fell about fifty-two percent as far as full fine-tuning's did.

What it saves

11:01 Ls dash l h g o lists the files with their sizes, in a readable unit, and leaves out the owner and the group. The fine-tuned model from episode three takes fifty-six megabytes, and LoRA's adapters about one point three. To keep another version of the model, fine-tuned for something else, we'd keep another file of adapters, next to a single copy of the pretrained model.

11:23 Side by side, the file we keep is about forty-three times smaller with LoRA, and AdamW's state is about forty-four times smaller, as it holds two numbers for each trained weight, and LoRA trains about two percent of them.

What it learns

11:37 Evaluate dot py measures three models: the pretrained one, full fine-tuning's from episode three, and LoRA's. On the held-out instructions' stories, the loss goes from two point six zero for the pretrained model to two point one three with full fine-tuning, and to two point three nine with LoRA. The requested words used go from about three percent to twenty-seven with full fine-tuning, and to eleven with LoRA, so LoRA's model follows instructions, but less often. On plain stories, its loss is two point three six, against two point three eight for full fine-tuning, so its loss on plain stories rose a little less from the pretrained model's two point two five.

12:14 Measured from the pretrained model, LoRA got about thirty-five percent of full fine-tuning's gain in the words used, and about forty-four percent of its fall in the loss on the instructions' stories, with about two percent of the weights trained, in ten minutes too.

12:29 With one instruction, rain, cake and jump, and one seed, the fine-tuned model's story uses jump, and LoRA's uses none of the three. One story is a small sample, and it fits the numbers, where LoRA's model uses about eleven percent of the requested words, and full fine-tuning's about twenty-seven.

At scale

12:48 At our size, full fine-tuning is cheap: ten minutes on a laptop, and a file of about sixty megabytes. So here, LoRA's savings matter little, and the gap in the requested words used matters more. A model of seven billion parameters takes fourteen gigabytes in bfloat16, two bytes a weight, and AdamW's two running averages, in float32, take fifty-six gigabytes more, before the gradients and the activations. With LoRA, the frozen weights get no optimiser state, and each fine-tuned version is a file of megabytes. So a server can keep one frozen copy of the model, and many adapters next to it, one per task or per customer. In an earlier run of mine, with the rank at thirty-two instead of eight, the model learned little more, so at our size, a larger rank doesn't close the gap with full fine-tuning.

13:35 Last, the docs get a page on LoRA, with the update and its scale, the code, and a table of what it costs and what it learns next to full fine-tuning, from an earlier run, along with the run at rank thirty-two. The docs' home page and their navigation link to it, after the page on PPO. And the README has the command to fine-tune with LoRA.

Wrap-up

13:56 So, today we fine-tuned our model with LoRA, which froze its weights and trained an update B A beside each linear layer of attention and the MLP, with B starting at zero. We fine-tuned on instructions for ten minutes, like in episode three. It trained about two point two percent of the weights, kept them in a file about forty-three times smaller than the model, and needed about forty-four times less of AdamW's state. In those ten minutes, it learned less than full fine-tuning: its stories use about eleven percent of the requested words, against twenty-seven, while its loss on plain stories rose a little less than full fine-tuning's. And merged into the weights, the update leaves a plain model, which runs as fast as before and loads like any other.

14:38 Let me know in the comments what you'd fine-tune a model on with LoRA, and which parts of it 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.