Fine-Tune Your Own Tiny-Llama on Custom Dataset

In the last video we looked at Tiny Llama which is a small language model. In this video I'll show you how to fine tune this model on your own data set. If you have watched the previous video, you will know that Tiny Llama is a very small model, so it's not really great for general language tasks. However, you can fine tune this model on very specific tasks and you can run this on edge devices. In this video we will look at how to fine tune this model on a very specific task. To fine tune this model we are going to be using a data set from Hugging face called Colors. However, you can just format your own data set. I'll walk you through a step by step process of how to do that before looking at the data set. I would like to say that this video is highly inspired from this blog post from Minyang Chain which is Tiny Llama colorist fine tuned with color data set. So I'm essentially taking his idea to replicate exactly the same behavior. Link to this blog post is going to be in the video description. Let's first understand the dataset before we look at fine tuning the model. So there are two columns. The first one is description. This column describes a given color. So for example you have pure black. Then here's a description of the color and the second column is colors, which is the corresponding hexadecimal code. Now we want to train or fine tune the tinylamo model so that it accepts the description as an input and then generates the corresponding hex code. Now you can do this task through prompting if you are working with a large language model by simply telling it to give you the hex code corresponding to a color description. However, in this case we want the model to just look at the color description. Generate the hex code without any other instructions. Here is an example. I took the Tiny Llama Chatbot, which is the instruct fine tune version, and I provided the color description and here's the output that it generated. This is basically text. However, in our case our fine tune model is going to generate the corresponding hexadecimal code without having any system instructions. I hope this clarifies what exactly we are trying to achieve here. So now let's look at the code. I'm running this on a free Google Colab and it has a pretty good inference speed as well. Now first we need to install some packages, so we will use Accelerate Bits and Bytes Transformer TRL for fine tuning the model. In this case we're going to be using the PEFT package from Hugging Face. And again, we're going to be adding a Lora adopter on top of the model. So we are not actually fine tuning the original model rather than the Lora adopters. If you want to learn more about Loras, I have a video on it. Link is going to be in the description. Next we need to import all the packages that we will need in order to fine tune the model. So we need Pytorch since we are downloading datasets from Hugging Face. That's why we need the load dataset function and dataset class. We are going to be using Lora. That's why we need to import the Lora configuration. Also the corresponding theft package. And for running the model, we'll use the auto model for Causal LM as well as the corresponding tokenizer. We are going to be using the supervised fine tuning trainer from TRL package from Hugging Face. Then we need to do some housekeeping. So first we need to provide the name of the data set that we want to use, then the corresponding base model. Now in this case, if you see I'm actually using the chat version of the model rather than the base model, you can do this on the base model as well. But just to keep the training time shorter, I'm using the already fine tuned version of the model that is trained I think on around 350 billion tokens. And then we need to define the output directory or where we are going to store this model. Next we need to properly format a data set so that we can start using this data set to train the model. Now the tiny Llama instruct version or chat version uses the chat ML format by default. And that's exactly what we're going to be using in our fine tune as well. So here is a function that you can use to properly format your data set. So you have the initial special tokens, then the input from the user and then the corresponding response. Okay, we're going to use this function later, but here's how I'm formatting my data set. So first we provide the data set id. Then we load the data set because it's available in a hugging face. We convert it to Panda data frame. Then we take both of the columns available in the data set. So one was description, the other one was colors. So the description becomes input to the model. That's where you expect the user input after these special tokens. And then the color column becomes the response which is after the special tokens for assistant. So if you were to format your own data set, it has to be formatted in exactly the same way. We convert that to a data set object again. Now, after formatting it in this specific format, we are creating yet another column which we are calling text. And this is critical. Whenever you're formatting your own data set, you need to apply the prompt template that you want to use and then assign it to a column called text. We convert it to dataset object again, right? And here is how we are processing the data. So we are calling this prepared train data function on our data id. And if I look at the data set, we have three columns. One is color description and then text, right? And There are around 34,000 examples in this data set. So if you look at an example here, this is the first data point or sample. So you have the color, the corresponding description, and this is the new column that we added. Right. So in the beginning you have the description which becomes the input, and the corresponding hexadecimal code which becomes the response from the model. Okay, so after this, we need to start setting up our model. In this case, we're using the chat version of Tiny Llama. So first we need to download and initialize the tokenizer corresponding to this model. The tokenizer is exactly the same like a llama 2 model. We're using the bits and bytes package. That's why we are setting all the corresponding configurations. Then we are loading or downloading the model itself. So if we call this function with the model id, it will download the tokenizer as well as the model and will set the corresponding bits and byte configurations correctly in here. And the return is going to be model and tokenizer. Next we need to set up the Lora configurations. So in this case I'm going with the default values. I have another video in which I go into more details on how to select different Lora adopters. So if you're interested, watch that video link is going to be in the description. Next, we need to set up the training parameters. So I'm not pushing the model to hugging face. That's why I'm defining an output model library or a directory in here. Since it's a smaller model, you can use a much bigger batch size. Even if you're running this On a free T4 GPU from Google Colab, we're using four steps for gradient accumulation. Here's the optimizer that we want to use. These are different things that you need to experiment with. Specifically, you want to experiment with the learning rate, which will control the speed of convergence of training. Now I'm Just running this for 150 steps, so it's not going through go through the whole data set multiple time. I think it's going to be a fraction of an epoch, so not even a full epoch. But if you want to run this for longer, I would recommend to either comment this out, just train it for a longer number of apex. Okay. After this we are going to generate or create a supervised fine tuning trainer object. So we provide the base model that we want to use, then the corresponding data. So in this case the data set has already been pre formatted for that specific prompt template. In my mixture of videos, I showed you how to actually provide the function which will format your data set. So if you're interested in doing it that way, have a look at that. We provide the LORA configurations in here. Since the data set is pre formatted by, that's why we need to tell the model which column to use for training. So we are specifically using the text column in this case. Then we provide the corresponding training arguments that we define up here. Corresponding tokenizer and the last thing that you need to come up with is the max sequence length. So I'm just keeping it to 1024, even though the data in the dataset is much shorter. So just to summarize, we're going to run this for 250 steps. And here is when we are training the model. So there was no validation data set, that's why it's only showing the training loss. And if you look at it, it actually decreases pretty nicely. Now there seems to be some fluctuations at the end, so that it's not really an indicator of overfitting yet, but we could potentially reduce the learning rate and then it's going to be much smoother. Another thing to pay attention to is that we ran this only for almost half an epoch, so not even a full epoch. Right. If you want to train your model, make sure that you run it at least for one epoch. Okay. So after this process it only trains the LORA adapters, not the actual model. So we need to merge both of them together. So, so I loaded the original model, then provided the path of the last checkpoint which was trained and loaded that. And then you want to merge and unload this model. So this is going to be the final model which has the LORA adopters merged with it. And the way you do it is you use this PEFT model class, then you need to provide the actual model that you want to merge the LORA with the corresponding path of the LORA adopters. And you get the final model. So if you look at the final merged model, you will actually see that it has the LORA adopters in here. Okay, so inference time. Let's see if the model actually learned anything. So for inference, we're going to be using this function called Generate responses. We get input from the user. That's going to be the prompt that the user provides. Then we use the formatted prompt function so that we can put this in the correct prompt template or prompt format. So here I put the definition of that function. So basically you receive a question, this will become user input with the corresponding special token. And then we're using special tokens for assistant response. So after this, the assistant or the trained model is supposed to generate a response for us. Now the rest is standard hugging face code. The only difference is here we are starting a counter basically to look at the start time and then we stop that. That's basically the time it took to make a prediction. Because this is a small model, so it's supposed to generate prediction at a really fast pace. Iwanona T4 GPU okay, and we look at both the output as well as the corresponding time it took to generate the output. Okay, so before testing it, here is another small helpful function which accepts hexadecimal code for a color and it will generate the corresponding color and show us. So we can actually see if the hexadecimal color belongs to the color that we provided or not. Okay, so the input prompt is a light orange color. The Response took only 0.35 seconds. So under a second, which is pretty amazing. And here is the hexadecimal code that the assistant generated. And for that color, here's the corresponding actual color generated by the function that we used. So we provide the response or the hexadecimal code generated by Tiny Llama and it's actually pretty close to light orange color. So it seems like our model is actually learning. This is pretty amazing. Now just to reiterate, you can easily do this with any large language model by instructing it to convert a given color description into hexadecimal code. However, in this case, the Tiny Llama model was able to learn from the data without any explicit instructions, which is pretty amazing. I'm actually excited about these small language models. I think we will see more of them during 2024, and this is one of the viable solution for running these devices on consumer hardware such as phones or edge devices. I hope you learned something new in this video and found it useful. Let me know if there are any questions. Thanks for watching and as always, see you in the next one.