Master Fine-Tuning Mistral AI Models with Official Mistral-FineTune Package

The team at Mistral AI just released an official Mistral fine tuning guide. It's a must read if you're interested in getting started with fine tuning Mistral models. They're calling this package Mistral Fine Tune, which is a lightweight code base that enables memory efficient and performant fine tuning of Mistral models. They're using Lora, which will add only 1 to 2% additional weights and that are going to be trained during the fine tuning process. This guide is extremely detailed and one thing I really like about this is they spent a lot of time on data preparation. I'll walk you through an example notebook that they have provided on how to get started with fine tuning Mistral models. But an important aspect of any fine tuning job is how to prepare your data set and how to correctly format it so the model can actually learn from your training data. So before looking at the fine tuning process, let's learn about how the data is supposed to be structured. If you are fine tuning Mistral model using Mistral Fine Tune. Now, for it to work, all the files need to be formatted in JSON ELF format. There are two different types of data sets that you can use depending on your training task. So if you're pre training, your data needs to be in a JSON L format where the key is going to be text and then you provide your unstructured text here. This is going to be useful in cases if you are just continuing pre training of a base model. Now if you are trying to fine tune a chat model or instruct model, then you need to arrange your data in the form of system message, user input and assistant response. Now they have divided this into two different types of instruction following data sets. One is a simple chat or instruct version. So in this case your data is going to be stored in a list with the key messages. So this is the base list. And then within these messages you need to assign a role. So if it's a user input, the role is going to be user. Then content is going to be what the user input is. Then assistant response is what the model is supposed to output. You can also include system messages if you want, and in that case the role is going to be system prompt. Now, during training, the loss is only computed on the assistant response because that is the predicted output that the model is supposed to learn. Now, in your data there might be some examples that you don't want to train on or they are not correctly formatted so you can assign it a weight of zero and the model will just ignore those during training. Okay. The second version of Instruct data set is to train the model for function calling.v3 version of 7B model from Mistral actually is fine tuned for function calling and it is pretty good at it. Now for function calling there's an additional role that you can use and that is the role of tool. So you can provide tools that the model is supposed to learn how to use those specific tools or function using this role. And in that case you will also need to provide examples of how those tools tool calls are made by the assistant. Okay, so if you're training it on your own data set, make sure that you follow the structure that is proposed here either for pre training or for instruct fine tuning. So again, your data is supposed to be in JSON L format. Okay, so they have provided an official Google Colab notebook that you can walk through and learn how to start fine tuning the Mistral models. I'm running this locally, so I'm going to walk you through how to set this up and what different sections of this notebook means. Now as you can see, I'm currently running a training job to fine tune the V3 of Mistral 7B. But let me walk you through the notebook of how I have set this up. Now I'm running this training job on a a 6000 GPU that has 48 gigabytes of VRAM and I'm using a virtual machine that is available through Mast Compute. If you want to use a similar system, link is going to be in the video description. Okay, so first we need to clone the repo and for that we're going to be using Git clone and then provide the URL of the Mistral fine tune package. When the cloning is done, you're going to see a file structure like this. So we have the Mistral fine tune package. This is exactly the same package that we were looking at on GitHub and within that you're going to see all different files. We're going to look at some of these files while we are going through the code. Now, now there is a requirements Txt file in this package. So I'm using that to basically install all the required packages that we will need during the training process. Next we need to download the model that we are going to be fine tuning. So in this case we are downloading the Mistral 7B V3 model and I believe this is the base version of the model. So this will directly download the model from Mistral AI website. Okay, so this will download the tar file which is Basically a compressed file. Now we need to uncompress this file. So I'm going to uncompress it in a folder called Mistral underscore Models. Now, instead of directly downloading from Mistral AI, you can also download the same files from Hugging Face Hub. Now, keep in mind that the Mistral models are now gated, so you will need to provide your Hugging Face API key or access token in order to download Models from Hugging Face. But I'm directly downloading from there, so we don't need to run this. Okay, so once the download is complete, if you look at the contents of this folder or directory, you can see that here are the files that are downloaded and actually downloaded the 7B instruct version. Okay, so next we're going to look at how to structure our data. And they actually have spent a lot of time in data preparation, which is pretty great because this is the step that is usually ignored in a lot of different packages that I have seen for training. And formatting your data correctly is going to be one of the significant steps that really determines the final output of the trained model. Now, for this job, we're going to be using the Ultra Chat 200,000K data set. But as I said, you can structure your own data by following the format that they have provided. So you need to format your data in a JSON L format and you need to have the system prompt user role and the assistant role in the JSON L format in order to structure your data. Okay, now we create another folder called Data. So this is basically the folder that is going to be created and then we move to that folder where we're going to download both the data set. So we will divide this into training and validation set and also make sure that the the format is correct. So they actually propose a step to validate the data formatting, and we are going to look at an example of how to do that. So I need to install a couple of additional packages in order to process this parquet file format. So we basically provided the path of the data set and we're going to load that into a data frame. So this is a Pandas data frame. Pandas is a great python package for data analysis. So you can look at elements that are present. So the way the data is structured is there's a prompt, then prompt id. We don't really need to worry about the prompt ID part, but then there is a part on messages. So let me show you what are the actual contents of this data frame. Now let's try to get A better understanding of the data itself. So we're going to randomly select an example out of the data set. So we are just looking at this specific example and there are three different values in here. So it seems like it's a multi turn conversation. The first one is a user role and that's the user message. Then there is an assistant role in here and that's basically what the assistant is supposed to generate. And this is followed by one more user message. So it seems like it's a multi turn conversation. There is also a prompt along with this. So I looked at the same prompt. So this one is Please write a well reasoned and informative article focusing on the evolving role of virtual reality technology. So this could be used as a system prompt. Now we're going to split the data set into two different parts. One is going to be the training part and the other one is going to be the validation part. And you always want to have a validation set which will be used to evaluate the performance of the model or like how good the learning is. So in this case we're going to use 95% of the data for training and then we drop those 95% examples and just keep the remaining 5% examples for validation. So we load the data in a form of PANDAS data frame, but we're going to store them in the form of JSON L format and that's why we're just using to JSON then providing the name. So this is the training section and the second one is the evaluation or validation section and we store them in JSON ELF format and the orientation is going to be in the form of records. So if you look here, we will get two different files and you can actually look at the contents of this file. So if we look here, here's what the data is structured as. So there's the prompt, then the promo, corresponding prompt ID and then user role and the assistant role. Okay, so now next they have provided some packages which actually checks whether the data is correctly formatted or not. So the examples that are not correctly formatted are going to be skipped and that is provided within this utils package. So there's this folder called utils and then you can see by using this form reformatdata file, this basically ensures that your data is correctly formatted. So this is a very important step that you want to run. So basically we're going to run this on the training data and you see there are some samples that are not correctly formatted and they're going to be skipped right now you want to do the exactly same step on the test data as well or the validation data. And this will also ensure that if there are any incorrect examples or incorrectly formatted examples, they are going to be discarded. And this is a critical step that you definitely want to run. Okay, so we looked at how to set up or how to format our training examples and validation examples. Now let's look at how to set up our training run. So you can use a YAML file with the configuration of your training data set as well as the training hyper parameter, which is pretty cool because unlike transformer package where you have to set up different values, here you just provide YAML file with configuration and then just run a training job on top of that YAML file. This setup is very similar to Xalot, which is another amazing library for fine tuning. If you haven't figured. If you haven't experimented with Xalot before, I'll highly recommend to check it out. If there is interest, I'll create a few tutorials on that. But the beauty of excelod is that you just need to provide a YAML file. So for example, this is how to create a data set. Then how do you set up model loading? Here is how you set up your Lora and you provide a single configuration file with these different hyper parameters. And then you just call the trainer on top of this and this will start training. Now the Mistral finetune package follows a very similar setup. So there is this example YAML file under the example folder. And if you go here, here's how you need to set up your data set. So first you need to provide the path where the training data is located. So this is the JSON file file that we created. Now if you're pre training a base model, you can also provide the path of the printed training data set, which we don't have in this case. And optionally you can also provide the path of your validation set. Then you will need to provide the path of the model that you want to fine tune. So we downloaded the model in this specific folder and we're just providing that path here. I believe you can directly also provide the repo ID on hugging face and it will download the model for you. Next you need to define the Lora parameters. In this case I'm just providing the rank, which is basically the rank of the Lora metrics that is going to be used. Now you can also provide the sequence length, what is going to be the batch size, max number of steps. Right. If you're using optum so there is basically for the optimization you will need to provide the learning rate, weight decay and so on and so forth. And the last couple of options is whether you want to merge the Lora adopters after fine tuning or not. In this case this is set to true, so it's not going to merge the Lora adapters. If you set it to false, then it will merge the lower adopters with the base model. And you want to provide a path of a folder that is going to be created during training. So you don't want to create a folder, but you just want to provide a name of the folder where the stored model is going to be created. Now we are going to create this configuration file and basically you just need to fill out these different parameters that are missing in the file and we're going to run this utils.investidate underscore data file to ensure that our data is structured correctly. The folder where we are saying we have stored the model is correct configured correctly. So this is basically a validation set or validation step that ensures that all the hyper parameters that you need for training the model are set up correctly. So when we run this, it says no error founds, data is correctly formatted. Right. So now we can start training. And you probably have noticed that we spend a lot of time on data formatting and just again confirmation of everything is set up correctly. This is one of the most important part that a lot of people don't pay attention to. But this will basically determine the accuracy of your final training. Okay, so we validated our data. The next step is to actually start our training job. Now we will need to set up our hardware. In this case I have a single gpu, so I'm just using that to run the training job. If you have a system with multi GPU multiple GPUs, the Mistral fine tune package actually supports multiple GPUs as well. So that's not going to be an issue. Now in order to run the actual training job, we will need to provide another configuration file. So either you can reuse the same configuration file that we had or you can set it up again. So what I did was I just copied the contents of this configuration file, which is a YAML file, into a variable called config and the code will generate another configuration file called example YAML. So if you look here, there's going to be a new example YAML file. This is basically exactly the same file that we had before. Now we're going to be running the training for only 300 steps and depending on the max sequence length in your training data. You also want to adjust this sequence length so they recommend either to use 32,000 or 800, depending on how large the examples are in your training data. One final note that this run directory is the directory or folder where your final model is going to be stored. But you want to make sure that you don't create this folder. Just provide the name of the folder that you want and whenever you run the trainer, it will create that folder. Now in my case you will see this, but since I ran the training job, so it actually created this folder for me which has the fine tune model in there. So once everything is set up, all you need to do is to use this single line of code to start a training job. Now this is very similar to something like auto train where you can use a single line of code to run the training. Axelot works in a very similar way, but we spend most of our time in just preparing and formatting a data set. Now when you run that, it's gonna start the training job. So initially it will just show you all the configurations that it is going to be using. So for example, here are the Lora Adopter rank and the other configurations. We only set the rank in this case. It will also show you other parameters that it's going to use. So for example, it will run it for 300 steps and in a process it will just create three different checkpoints. Okay, so it will run the training job. You can actually see that the loss decreases pretty nicely throughout the training job. And we also have the validation loss that we can track to see whether the model is overfitting or not. Now this training job took an hour and a half on this A6000 and we were looking at about 200,000 examples. So I think a pretty good speed at which we were able to fine tune the model. So let's look at some of the values in here. So the Training loss is 0.08 and the validation loss is around 0.10. So not bad at all. I think it's able to train the model. So this is pretty good. Now once you train the model, the next step is definitely want to see whether you can use this for inference. So can you use the model to generate text? They recommend to use another package called Mistral Inference. This is their dedicated package for running inference jobs or production jobs. So we are installing that package here. Okay, so we are going to load all the packages that we will need. Now keep in mind that we only trained The Lora adopters. So we will actually need to merge that Lora adopter with the model. So here is how you do it using the Mistral inference package. So first we will need to load the tokenizer and they recommend to use the V3 tokenizer moving forward for all the install models. Next we were going to load the base model. So this was the model that we actually downloaded initially and we were to be merging our Lora with this specific model. Now if you go to this fine tuned model, you will see that it has multiple checkpoints. So this is checkpoint two, this is checkpoint one and I believe there is a checkpoint three as well. So you can select any one of them. In this case I'm just selecting a third checkpoint and we will just load that checkpoint and merge it with the model. So basically we are training, we are loading our train lora and then merging it back with the base model. Okay, so once you do that, we'll have a fine or fine tuned model. Now we are loading this in 16 bit precision and we can run our prediction job. So to run our prediction job we will use this chart completion request package. Then we'll provide the user input. So the input is explain machine learning to me in a nutshell. Then we take that, pass it to the tokenizer to create tokens which the model will be able to understand. And we're using the V3 tokenizer in this case. Then we run that through the model. So we're asking the model to just generate a max of 64 tokens. If you want a more detailed response, you want to set this up to a higher value, temperature is set to zero and we also need to provide an end of sequence token. So this will generate output tokens for us. Then we take those output tokens, run them through the tokenizer again in the reverse order so that it will generate English text and we print that. So the response is. Machine learning is a subject of artificial intelligence that involves the use of algorithms to learn from data and make predictions or decisions without being explicitly programmed. So a pretty nice answer for this fine tune model. So this was an overview of how to fine tune a Mistral model using this new Mistral fine tune package. Now you probably realize that we spend a lot of time on how do you format data, how do you verify that the data formatting is done correctly? And this is one of the most critical part that really make or break your training run. So make sure to pay close attention to that. I create a lot of content around fine tuning or training LLMs as well as how to design Rack systems different components. So if that's something you are interested in, make sure to subscribe to the channel. Also, push that bell notification button so you never miss another video. I hope you found this video useful. Thanks for watching and as always, see you in the next one.