Fine-tune Mixtral 8x7B (MoE) on Custom Data - Step by Step Guide

This is probably one of the most requested video on my YouTube channel currently. Today we are going to be looking at how to fine tune a Mixtral 87B, which is the mixture of experts model from Mistral AI on your own data set. The idea behind mixture of Expert or MOE has gotten a lot of traction recently because it's speculated that GPT4 is potentially a mixture of experts. The MOE or Mixture of experts from Mistral AI is also really good for its size and it's able to beat GPT 3.5 on a whole bunch of benchmarks. So in this video I'll show you how to fine tune this model on your own data set. In this video we are not only going to be looking at the code, but I'll also talk about some of the practical things but that you need to consider when you're fine tuning any LLM or a machine learning model for that matter. Now, first and foremost you will need around 60 to 65 gigabytes of VRAM in order to fine tune this model, even if you're using S, Lora or Qlora. Thanks to my friends at Sagex AI, I got access to an H100, so I thought I'll just use that. So let's get started. First we need to install all the required packages. In this case we're going to be installing transformers TRL for training the model, accelerate Pytorch bits and bytes. And for training we're going to be using a hugging face dataset. But I'll show you how to format your own dataset so that you can use it if you want to fine tune this model. And we are going to also use Flash. Attention. If you want to train this model or fine tune this model, you probably need four T4 GPUs which comes out to be around 64 gigabytes of VRAM. Next we need a data set to train the model. In this case again, I'm using the mosaic ML instruct v3 dataset. This is the same dataset that I used in another video where I showed you how to fine tune the original Mistral 7B model. Let's have a quick look at the data set. The data set has two different splits. One is the train and the other one is test split. There are three columns in here. One is the prompt. So the prompt contains a system message. Then the actual user input and the corresponding response is another column. This data set has data from multiple sources. So you can see this. Here is Dall E. Then Computation math. This is chain of thought GSM8K. So the third column shows where a particular example is coming from. I discussed this data set in a lot more details in my previous video. So if you're interested, I'll highly recommend to watch that. Let's have a closer look at the data. There are about 56,000 samples or examples in the train set and around 6,800 examples in the test set. Again, there are three columns. Prompt, response and source. Most of the data in this data set is coming from this one specific data set. I think this is Dolly Harmful, harmless. This is coming from Anthropic. Now, when I filtered it for this specific data set, we have 34,000 examples in the train set and around 4,700 in the test set. So this subset is the data that we are going to be using to fine tune our model. Next, in order to fine tune the Mixtral A7B, we need to format our data and bring it into a single column. We need to format our data in a very specific format which we're going to be calling the prompt template. Now, if you look at an example from the data set, here is how it's formatted. So you have the prompt which contains a system instruction. Then what the user wants from the model. So basically a question and then special tokens response. And here's the response that the user is supposed to get from the model. But the mixture 87b instruct version follows a different format and the format looks something like this. So you have these special tokens in the beginning, Then you have the system message user input, then another special token which indicates the end of the user input, then the model is going to generate a response and then there is a special token which indicates the end of the model response. If you are fine tuning the instruct version of the Mixtral 87B, you need to provide your data in this format. But if you are fine tuning the base model, which is the next word prediction model, then you have the freedom and flexibility to define your own prompt template. Now, in my case, even though I'm fine tuning the base version of Mistral 87B, I will still follow the prompt template that Mistral recommends. Alright, so this function will basically reformat the initial data into the prompt template that we want. Now, to make it a bit more challenging, we will also rearrange our data. So instead of just providing a system message or system instruction, a question from the user and getting a response from the model, we will rearrange it. So our system message is going to be this, which states Use the provided input to create an instruction that could have been used to generate the response with an LLM. So basically what we are doing is we are providing text as an input and we are asking the LLM to generate certain questions that could have been asked by the user to get the text that we are providing to the model. Here is an example of reformatted data. So this is the initial instruction that we are giving to the user. Then here's the text that we are providing to the model. And the model is supposed to generate a question based on the text that we provided. Next we need to load our base model. But before that, as I mentioned in my previous video, this notebook is based on the amazing work done by the folks at AI Makerspace. Do check out their channel link is going to be in the description. As I mentioned, we are using the base version, so this is not the instruction version. Next we will set some configuration. We're going to load the model in four bits, but for compute we're going to be still keep it in 16 bits. Okay, so this comment doesn't apply here because we're using the base model. Although for fine tuning the base model correctly, we'll need a lot more data compared to when we are fine tuning the instruct version of the model. Next we will load the actual model. So we provide the model ID. We want to use all the GPUs available on the system and we're making use of a Flash attention too. So if your GPU supports Flash attention, make sure to enable it. That will expedite your training process. After that we are also loading the tokenizer. We are setting the end of sequence or end of sentence tokens. Now not all the sentences are going to have the same length, so we need to also pad the ones that don't have the max number of tokens or max sequence net. The neural network or transformer architecture in this case is performing mathematical computation. So you want your data being to be in the metrics form, but for that you need to have each and every row have the exact same dimension or the number of columns. So let's say if your max sequence length is 200 tokens, but there are examples which have only 150 tokens, then you want to pad those examples so that their length is equal to 200 and the matrix computation is going to be a lot more efficient. Now the question is how do you determine the max sequence length? So in order to explain that, we're going to look at an example later in the video. So before fine tuning the model, let's look at what kind of responses we get out of this base model. So I'm going to be using this function which basically takes a prompt from the user and the model generate responses. Okay, so let's look at an example. This is going to be my input prompt. And the prompt is use the provided input to create an instruction that could have been used to generate the response within LLM. And the text that we're providing is this. So there are more than 12,000 species of grass. The most common is Kentucky bluegrass because it grows quickly, easily and it is soft to touch. So this is basically the prompt that we are providing and we are asking the model to generate potential questions that could have been asked by a user from an LLM. Okay, so when I ran this prompt through the model, here's the response that it generated, I simply reformatted so that it's easily readable. For some reason it added a couple of other special tokens. But anyways, here's the original prompt that we provided and after that the model, since it's a base model and it just does sentence completion or next word prediction. So it reiterated my prompt and after that it simply started creating a whole bunch of random special tokens. This is the expected behavior from a next word prediction model. It's not really following the instructions and that's absolutely fine. Let's see if we can fine tune this model so that it can start following these instructions correctly. Before looking at the fine tuning process, let's again talk about the tokenization and the max sequence length that your fine tune model is going to support. This is very important because depending on how big your max sequence length is, the model will take that much longer because it will need a lot more computation. Mixtral87b out of the door supports 32,000 tokens. But if you are working on a specific application which doesn't need that big of a sequence, then you want to fine tune it with a much shorter sequence then. And let me show you how you do that. So this is a code snippet that I borrowed from Brave AI. I think they have an example notebook. I'll put a link to that in the description. So basically what we are doing is we have our Create prompt function that we use to reformat our prompt. So I am reformatting all my training data as well as test data and then passing thread to the tokenizer of mixtalk87b so we will get a tokenized data set. Now after that, using this function, we are Looking at the distribution of the sequence length of each and every sample. So here is the frequency distribution. Most of the concentration of this histogram is in the lower part. So that means that most of the examples that you have in the dataset, they have a sequence length of probably somewhere around 7 or 800. Here I just looked at a relatively zoomed in plot of the same data and that most of the sequences or examples that you have, they are less than 1200 tokens. If you're fine tuning a task specific LLM, just look at the distribution of length of tokens of each example and make sure that you select a max sequence length so that it covers most of the examples. There might be some outliers, you can just throw them out. Completely ignore those. This will save you a lot of computation. Next, we're going to quickly look at the architecture of Mistral 87B. Just pay close attention to these layers, the linear layers that you see in here. These are the ones that we're going to be adding Extra adopters using LoRa. Next, let's talk about our LORA configuration. So these are different modules or layers that we want to attach our these LORA adapters to. So if you go back, you will see these model or block names and these are the ones that we're going to be targeting. Now, I have covered LORA in a previous video. I'll put a link to that video. After setting the specific target layers that we want to address, you can play around with these different parameters. For example, you can define the LORA dropout. So basically this will regularize the model and potentially reduce overfitting of the model. Right. So after setting those up, we apply the LORA adopters to our model and we're going to be just training these LORA adapters. So here, using this function, we can look at out of all the parameters that are available, there are a lot we are just training 0, 24% of all the parameters that are in the model. So essentially we are just fine tuning or updating the weights of these LORA adapters, not really playing around with the weights of the model. If you want to fine tune the whole model, that will take a lot of compute, a lot of time and a lot of resources. And that is the beauty of loram. Okay, now let's look at the model again after attaching our lora. So if you see here, this is one of the linear layers and we added the LORA dropout, then the LORA adopters to that. Then we took the second linear Layer. And again we're going to add both the Lora dropout and Lora adapters. So if you notice, we're only targeting the linear layers. We are not playing around with the non linear layers or even the router module in there. Okay, next we're going to set the hyper parameters for training. So first we need to check if you have multiple GPUs on the system. If you have multiple GPUs, we will enable parallelization of model training. All the training is being performed to the amazing TLR package from Hugging Face. Now we need to define the output directory where we want to store the model. Then in order to train the model, you have two options. Either you can define the number of epochs that you want, or you can define the maximum steps that you want. So for this quick experimentation, we're going to set it to 250 steps with a huge batch size of 32. Then for we're going to look at the results every 10 steps. And I'm setting a relatively lower learning rate here. And I want to have a smooth training process. The convergence will take longer, but hopefully the loss is not going to shoot up over time. Now these are different parameters that you definitely want to play around because these will control your training process. Okay, so with all that set, we need to next set our trainer. So I'm using the supervised fine tuning trainer from TRL package. Now again, max sequence length. I basically looked at this plot and I think if I set it to 1024, it covers majority of the samples that I have in the training set. Now you can manually discard the samples that have a sequence greater than whatever you set in here. But in this specific example, I'm not doing that. Another thing that you will notice here is formatting function. So I did not not reformat my training data set. What I'm doing is I'm going to reformat that on the fly using my create prompt template function. And we are doing that both for the training as well as the test data set. Okay, so on my H100, thanks to my friends at Sagex AI, it took around 85 minutes to complete 250 steps. If you look at the training as well as validation loss, there's a pretty nice gradual decrease in training as well as validation loss. So we don't really see any sign of overfitting in here. That's a good indicator. It also means that I could have potentially trained it for longer to get better performance out of the model. All right, so after training the model, we can Just just store the weights in a local folder. You can also push this model to hugging face using this command. So you will need to install the hugging face CL for the notebook. You can use this. This will ask you for your credentials and after that you can just use the trainer push command, provide the repo ID where you want to push your model and it will upload your model there. Now we only train the Lora adopters, so in order to use this model we will need to merge both the original base model as well as the adopters that we trained. But let's see how good this model actually is. Now here I'm using the same function to generate a response. We will use exactly the same prompt, but in this case the model is the merged model. Here's the response that I got from the model. I just formatted it for readability. So initially we have our input and the output is this. So it says use this input and we expect the model to generate questions. The first question is how many species of grass are there? What is the most common grass? And what grass would you find in a very dry area? Now these are definitely relevant to the text that we have in here. Now, for some reason the model actually repeated the same question again, but at least it's following instructions and it's generating questions. If we train this model for much longer, we will start getting much better responses. And the reason is that we were using only 250 steps with the batch size of 32. So the model basically looked at only 8,000 examples. It hasn't even looked at all the data set yet. That's why if you have a huge data set, you want to train this model at least for one or two epochs. In some cases probably more than that depending on how large your data set is. Okay, so this was a very quick tutorial on a step by step process of how to train the newly released Mistral 87B, a mixture of expert model from Mistral AI. Let me know in the comment section below if there are any questions or if there are any other specific topics you would like me to cover in my future videos. I hope you found this video useful. Consider liking the video and subscribe to the channel for more content like this. Thanks for watching and as always, see you in the next one.