Mistral: Easiest Way to Fine-Tune on Custom Data

The real power of open source large language models comes from the ability to be fine tuned on task specific datasets. Mistral 7B is one of the best option if you're looking to fine tune a small large language model. In this video I will show you how to fine tune a Mistral 7B on your custom dataset. I'll also show you how to correctly format your dataset for fine tuning. And at the end of the video I'll show you an alternative that you can use to fine tune large language models on your own dataset if you don't have access to powerful GPUs. So let's get started. I'm using the Pro version of Google Colab. I'm using an A100 GPU that has 40 gigabytes of VRAM. Later in the video we are going to look at Gradient, which is a platform that you can use to fine tune and serve your LLMs. And they are the sponsor of today's video. The credit for this notebook goes to the AI Makerspace YouTube channel. They have an awesome channel and I would recommend everybody to check it out. So first and foremost we need to install all the required packages. So this includes Transformers, TRL Accelerate, Pytorch, Bits and Bytes. We will use PIF for training the model and we're going to be using a data set from Hugging Face. But I'm going to show you how to format your own data set. Once the installation is complete, we will be loading the InSTRUCTV3 data set from Mosaic ML. They were the creators of the MPT models. These were some of the early open source large language models. So let's try to understand how the data set is structured. Each of the data set has three columns. The first column is the prompt, that's the user input. Then there is a model response and then there is a source column. And we are going to look at this to learn how the data is formatted. Now there are two different parts or splits. So there is a training data and there is a separate test data in terms of the composition of the data set. This data set is basically a combination of nine different other data sets that are put together in here and they contain databricks. Dolly 15k which was combined with the anthropic, helpful and harmless data set. Then there are a few other data sets which are listed in here. So this include computation math, there is this well known GSM 8K and then some screen, FD and spider data set. But we are going to be simply interested in the Dolly section of the data set for this video. Now, let's understand the structure of the data. So first we downloaded the data set and we are calling it Instruct tuned data set. And if you look at the dataset dictionary, there are two components. There's a train split and test split, and each of them contains three different columns. Now, the source column will define where the data is coming from. So it's going to be one of those nine distinct data set. And we're going to be combining these two together in a specific prompt template that works with Mistral 7B. First, we're going to simply filter this data, and we just want to look at the dolly harmful and harmless data set. So the way we do it is we are using this lambda function which looks at the the source column for each sample. And if it belongs to this dolly harmful and harmless data set, then only that is included in the final data. For all the rest of the eight values, they will be simply ignored. So as a result, now you can see that our train set has 34,000 examples, whereas our test set has 4,700 examples, if you look at the original unfiltered data set. So it had around 57,000 examples for the train set and around 6,800 examples for the test set. We are not going to be training our model on all of these 34,000 examples because that is going to take a while and it's going to cost us a lot of money. But we're going to just use a subset of this data set, and here's how you do it. So I'm simply selecting the first thousand examples out of this train set and then the first 200 examples out of the test set. Now, if you look here, we have thousand examples in the train set and 200 examples in the test set. However, I don't want to use the original format that it's structured in. So let's go back and look at the data again. So here is an example from the data set. So at the top you have a system prompt or user input which states, below is an instruction that describes a task. Write a response that appropriately completes the request. Then the instruction is a question, and the model is generating a response based on the question. However, the way we want to do it is we want to provide this as an input, and then we want a model to generate an appropriate question for the text that we provided. So this is going to be the input to the model, and then we want the model to generate a question which is going to be the Response of the model. Now, with that in mind, we are going to create a single column that is going to combine both the prompt and the model response. And this is exactly how you want to structure your own data set as well. So the prompt template that we want to use is something like this. So there's going to be an instruction which will state use the provided input to create an instruction that could have been used to generate the response within LLM. Then there's going to be a special token for input. This is where we are going to provide text. So let's say this is the answer that the model is going to use to generate a question. And then there's going to be another special token followed by the question that was used to generate the response. So this is how you want to structure your own data sets. So here's a Python function that does exactly the same thing. So here's how it works. We have the beginning of the sentence token, then the original system message or system prompt that is coming from the data set. So if you go back and look at all the examples that belongs to this dolly data set, they have this special system instruction or system message in the start or beginning of each and every prompt. So what we want to do is we want to replace the original system message by this new system message which states that use the provided input to create an instruction that could have been used to generate the response within LLM. So imagine that you get a sample from the data set. Then we are replacing the original system message with this new system message. And the way we do it is we replace the original system message with an empty string, the special token for instruction with the empty string, as well as the special token for response with empty string. And then we take the response from the input training sample and put that as an input. And now we're creating a full prompt. So the way it works is that initially you define an empty string, then you put the beginning of the sentence special token. So this is going to be the special token that is going to be put in there. Then we are going to append it with this special token for instruction. Then next we are going to add the system message. This is our new system message. Then we add special token for inputs, and then we add our input. So this is basically the previous response from the model. This. Then again, special tokens for response. We add the response that we got in here and we add the end of the sentence token. So here we can look at an example. So we pass the first training sample as an Input and the output is this. So you see the instruction, the instruction has been updated to our new system message. So this is the one. Use the provided input to create an instruction that could have been used to generate the response with an LLM. So this looks correct. Then we have the input and the input is actually a small paragraph and the response in this case is the original question that was asked to generate the text that is provided as input in here. So it seems like this is working. So we just looked at an example of how the output is going to look like from the Create prompt function. However, we haven't applied that to the data set yet. You can actually transform the whole data set using the Python map function. So all you need to do is just provide the Create prompt function in here and and this will map each and every training, as well as test examples in the data set to this new format. However, instead of doing this manually in here, we are going to do this as a part of the training process. So I'm going to show you that in a bit. Okay, now we are all set for training our model after the data pre processing. Now in this case we are going to be using the transformer package to do the training. Now, even though we are running this on a 100 GPU which has 40 gigabytes of VRAM, if we were to fine tune the model in full precision, which is 32 bit, we will need a lot more VRAM. That's why we are going to bring it down to 4 bit. So we are going to load the model in 4 bits. That means we will need only 25% of the compute that is going to be needed for something like 32 bit precision. However, for actual model weight updates, we are going to upscale this to 16 bits just to ensure that we do not lose a lot of information. Even with this, I saw that the VRAM consumption goes up to 32 gigabytes when we are training the model. Not everybody has access to something like a 100. Now keep in mind, I'm not loading the model and updating them in different shards. You can do that as well. I have a video on that. So in that case you could potentially use something like a T4 GPU, but the performance of that model may not be as good as doing it in 16 bit updates. Now, if you are resource constrained when it comes to training LLMs, that's where the sponsor of today's video comes in. Gradient is a platform that lets you fine tune your own large language models. Not only that, but it lets you serve them to an API so you can access them anywhere you want. All you need to do is provide you a data set and they will take care of the infrastructure part. They have a number of large language models to choose from, including llama2 nocermese, llama2. They have a very powerful Python SDK. Simply choose the base model that you want to fine tune on your own data set, set up all the parameters for training, provide your DataSet in a JSON format and let the system take care of everything else. Another great thing about them is that they have an embedding API which serves an open source embedding model that you can use in your own RAG application. And it has integration both with LangChain as well as Llama Index. So with a single API call you can bring in your custom fine tuned model into Lagchain in Llama Index. Do check them out, I'll be making a lot more content on them. Now back to the video. Next we want to load the model that we want to fine tune. So for that we're going to be using the auto model for causal LM and we'll be using a pre trained instruct fine tune Mistral 7B. Now keep in mind, I'm not using the base version of Mistral, but rather I'm using an instruct fine tune version of Mistral. And the reason is that the base models are simple next word prediction models. So in order to fine tune them on an instruct data set the way we have configured, we will need a lot bigger data set than thousand examples for it to work perfectly. The instruct fine tune version is already fine tuned on question answer pairs, so it's a lot easier to fine tune that with new knowledge. And we are loading this on our gpu. So we let it decide to map the GPU that it want to use. There's a single A100 GPU so we are not going to run into any issues. And now we load the model in four bit. So apart from the model, we also need the corresponding tokenizer and that's how we are loading our tokenizer. Okay, so before fine tuning the model with our own data set, let's look at how the base model that we just loaded is going to respond to our new prompt template. So here's the prompt. So we put the instruction in the beginning which is use the provided input to create an instruction that could have been used to generate the response with an LLM. So here is our input text. There are more than 12,000 species of grass, and so on and so forth. Now, in terms of the response, we were expecting a single line which is going to be a question, but it generated a pretty dated response. Even though we provide this specific instruction to the model to just give us an instruction that could have been used to degenerate the response. So you can already see the model is good, but it's not doing what exactly we want it to do. And that's where the fine tuning is going to help. So next we are going to look at the actual training part. Now for training we're going to be using LORA that is implemented as a part of the PEFT library from Hugging Face. Now, to understand the concept of Lora, let's look at this simple diagram. The original model weights that we have when we load the Mistral 7B, there are a whole bunch of weights that are simply redundant. And LORA uses that to actually reduce the size of trainable weights that we need. LORA stands for Low Rank Adaptation of Large Language Models. This was the paper that introduced this concept of LORA to keep the description very brief. So instead of updating the original weights of the model during the fine tuning process, we actually append weights that are much smaller in number and we update those during the fine tuning process and later on we simply merge both of them together to get the updated or merged fine tuned weights. And this is a technique that helps you a lot in saving the amount of VRAM that we will usually need in order to fine tune these large language models. So here we are just defining all the configurations for the LoRa. Then we need to apply those configurations to our model. So we're going to be using the PREPARE model for K bit training. This is a function within the PEFF package from Hugging Face. Right? And simply update those. So we have a new model with the LORA applied, but we haven't trained the LORA yet. So we need to set a few more hyper parameter for training. And again, as I said in the beginning of the video, this is based on a pretty awesome notebook from the AI maker space. So do check out the channel. So here we provide the output directory where the train model is going to be stored. Now, when you are doing training, you can either use the number of effects that you want to run or the number of steps. So I'm going to explain this in a little bit. Now you need to select the batch size. So this is basically the number of examples that is going to be used at once to train the model. Then you have some other parameters, but the most important parameters that I would recommend you to keep track of is the learning rate. This will control the conversions or speed of training. Then you also want to be very careful about the number of epochs or the number of steps that you are going to train the model with. Because if you train it for too long on a small data set, you will probably overfit the model. So let me explain the difference between epochs and steps. So epochs is basically one iteration where you process the whole training data set through it once. So let's say if you are using a training data set of thousand examples with a batch size of four, that means you will take a total of 250 steps so that you can have a one iteration of the whole data set through the model and that is going to be your one epoch. Now the number of steps are basically how many batches you are passing through the model. So we are setting it to a small number in here just to save on, save on compute cost. So either you can go with a full apex or a sub apex which is going to be controlled by the max number of steps that you want the model to go through. Then at each 20 steps we are going to use the test set that is available to us and evaluate the performance of the model. All right, so this is the part where the actual training is taking place. So we're going to be using the supervised fine tune trainer. From Hugging Face. In here we need to provide our model. So keep in mind we added the LORA adopters to the model. Then we're providing the PEFT configurations that we already defined. This is the maximum sequence length. Even though the menstrual 7B has a much larger sequence length, we are just limiting it to smaller sequence length because the both the input that we are providing as well as the response that is supposed to generate is going to be within 2000 tokens. Then we provide our tokenizer. Now here's the most important part that actually need to consider. As I said in the beginning, you can use the map function to map your data set to this new format, or you can provide this function to the formatting function in the SFT trainer. So that will format the prompt template on the fly when it's processing the data through the SFT trainer. And then we provide our training set and our test set. Now after that you can simply run the training process Here you can see that both the training as well as the validation loss are going down. So that's a good indicator that the model is actually training and you can see that this ran for 0.4 apex. So one the hundred steps that we define is probably around 0.4 of 250 steps that is going to take to create a single epoch. And at the end we're just storing the model to a local directory. Now you can push this model to GitHub if you want to use it later on. So the way you do it is you provide your GitHub credentials. If you run this block, it will ask you for a token. You can provide that then provide the repo ID where you want to push this. Now it only pushes the LORA adopters, so you will actually need to merge those to your original model if you want to use it again. So let's test the model. So here's a function that will generate a response and in this case I provided both the prompt as well as the model. So let me show you. This is the merged model that we're using. The text or the input describes how to make guacamole and the model response is how to make guacamole. Right. So it is actually working pretty fine with a relatively smaller data set that we use for training. So that's how you fine tune a Mistral 7B model on your own data set. Now just to recap, you need to be careful about the pre processing of your own own data set. That is I think the most important part. The rest of the parameters that we saw during this video are pretty much the same that you can use now. You can also play around with the learning rate if the training is too slow. You can also play around with the number of Apex. That depends on how big of a data set you have, because it's going to take quite a while if you have a large amount of data set. I hope you found this video useful. Let me know if you want me to make content like these and go into a lot more details of the training process. I would love to do that. I already have some videos on how to format datasets using GPT4, so I'll put a link to that video in the description. Thanks for watching and as always, see you in the next one. Thanks.