Custom Fine-tuning 30x Faster on T4 GPUs with UnSloth AI

If you are working on fine tuning open source large language models, this is a tool you definitely want to check out. There's a new AI company called Unsloth AI and they just introduced a new package which can make training time 30 times faster. That means training a model on the Alpaca dataset will take only three hours compared to 85 hours. Now we have quite a few open source projects for increasing the speed of inference such as Vllm, but this is one of the few projects that is focused on the training aspect. So first we will look at how they are able to achieve the performance increase and then later in the video I'll show you how to use this in your own projects. So according to the authors you will need 60% less memory usage which will allow six times larger batch sizes. Now everything is done through software changes, so you don't need any new hardware. And the great thing is that it supports Nvidia, intel and AMD GPUs. Now how they are able to achieve this performance? Now how are they able to achieve this? Well, the authors have hand implemented AutoGrid in PyTorch and they have also applied some other mathematical tricks. Now a major part of it is that they rewrote Auto all the kernels in OpenAI's Triton language and they also make use of the Flash attention mechanism. They are offering three different plans. There is a free plan, then Unslot Pro and the last one is Unslot Max. With the Free plan you still get 2 times faster training with 40% less memory usage and it supports both llama 1 and llama 2. Free plan is limited to a single GPU's usage. For the Pro plan as well as for the Max plan you get up to 32 times faster training, which is amazing. So let's look at some performance numbers before we look at the code on how to use this package in your own fine tuning projects. They have given performance analysis for two different data sets. The first one is Leon chip two 210,000 examples. So they ran the training for one epoch on two T4 GPUs. In this case the open source implementation or the Free tier is able to do training almost five times faster than the Hugging Face Transformer implementation, which is pretty amazing. And the Max version goes up to 32 times faster, which is pretty awesome. Now keep in mind that for the Open or free tier they're just using a single gpu, which I think makes it even more impressive. Performance doesn't scale with data sets. So for example, for A larger data set of almost half a million examples. The free tier gives you 1.6 times faster training, which is still very impressive on the hugging face repo. They also provide the peak memory usage for different data sets. Now, if you look here On a single T4 GPU, if you look at the alpaca dataset, so hugging face Transformer takes around 7.3 gigabytes, whereas the free tier takes around 6.8 gigabytes for the Alpaca, 50,000 examples. Now interestingly enough, is for much larger data sets such as the Slim Orca dataset, which has around 518,000 examples. The memory usage for the free tier is around 7.7 GHz gigabytes, which means that you can easily train this on a free tier of Google Colab. So now let me show you a quick example of how to use this package and get faster training on your own projects. This is a notebook that is provided by the authors of the package. So first you need to install the unstalled package from GitHub. Then we are importing this special class called fast llama model. From the onsloth package. We're also importing torch setting up the max sequence length. The data types is set to none. So this is basically for auto detection based on the system. It can automatically set it to float 16 or B float 16 and then we're loading the model in four bits. Now, currently the unslot package only supports llama 1 and llama 2 base models. So something like Mistral is not supported yet. Now in this case we're going to simply load the corresponding model as well as the corresponding tokenizer. Here we are setting the different modules of the model that we want to train. Set of different hyperparameters for Lora and then we set the max sequence length. So I think in this case we set that to 2048. I think at the moment it only supports the alpaca prompt template or format. So here's the alpaca prompt template. So basically they define this formatting prompt functions to reformat your data set into the alpaca format. For this example, they are using the alpaca cleaned data set from hugging face. Now let's look at the training data itself. This has around 52,000 examples. And then basically we are using the prompt formatting function that we defined above to reformat the whole data set. In the original data set you have three different columns. One is the output, second one is the input and the last one is the instruction. However, here we are reformatting those into instruction, input and response. Now the only change that you will see between this and other code examples which use the transformer package for fine tuning model is the way the model is loaded. So here we're using the fast Llama model class to load the model. That's the only change. So for the rest of the training we're going to be using the SFT trainer, which is the supervised fine tuning trainer for training the model. And the training arguments are going to be coming from the transformer package. So here we provide our model and the corresponding data set. Now the data set is converted into a single column. So here's the text column that is being defined and it's going to be in this format. So we are telling the SFT trainer to use that column to retrieve the data. Then set up the max sequence length. Now for the training arguments we are setting up different parameters. The most important is the max steps. So we are just running the model training only for 60 steps, so not even full epoch in this case. And we are setting a small learning rate to begin with. Then we are using a linear learning rate scheduler which will linearly decrease the learning rate with different steps. And for optimizer we're using the item optimizer. Now here's the fun part. So, beginning of the training we have around 15 gigabytes of VRAM. And for the training process only around 4 gigabytes is reserved, which is pretty awesome. So we kick off the training process. Now the 60 steps took around eight minutes, which is I think pretty impressive. And the training loss does decrease gradually. There are a couple of places where it kind of jumped, but it starts decreasing again. So I think we probably want to look at the batch size as well as the learning rate a bit more. But overall there is a gradual decrease in the training process, which means that we are learning. It's definitely nice to have a separate validation sets to ensure that there's no overfitting. But for this example they didn't provide any validation set. This is probably something you want to definitely look at. Now the training is complete. So now let's look at some of the statistics. So the peak reserved memory for training was around 5.27 gigabytes after of the 15 gigabytes that is available. So that means that at peak reserved memory we were only using around 35% of the available VRAM, which is very impressive. Now keep in mind that we are running the training on a free T4 GPU from Google Colab. Now to be absolutely clear, you can actually fine tune a 7 billion model on a T4 GPU without the need of this unslot package. However, in that case you will need to use the sharded version of the model rather than loading the full model in memory. So basically in very simple terms, you will divide the model into different shards and you will load one shard at a time. So not the whole model is going to be in the memory, but in this case you can keep the whole model in memory and the training is a lot faster. The training works, but I also wanted to see whether you can make predictions on the trained model or not. So here's the code that I wrote. So first I saved the model to the disk on my Google Colab notebook. Then I wrote this function to get prediction out of that model and it actually works. So I asked the question what is the meaning of life? And the model responded what is the meaning of life? Is it to be happy? Is it to be successful? Is it to have a good family? Is it to make a difference in the world? And then I think there was some weird token in there. Then it goes on to say there is no one size fits all answer to the question as the meaning of life can vary greatly depending on the individual's beliefs, values and priorities. So the model actually works, which is pretty amazing. In this video we tried the free tier of the unslot package and I think the paid tiers will definitely have a lot improved performance and speed. So if you're fine tuning a large language model, I would definitely recommend to use this for training and then you can use a package like VLLM for inference. And serving a combination of both of these packages can definitely reduce both your training as well as inference time and hence the cost associated with fine tuning and serving your models. I'm going to put a link to the Google Colab in the video description. This is based on the provided notebook. I just added the inference part in the notebook, so do check it out, Let me know what you think and what kind of performance you're getting out of this package. Thanks for watching and as always, see you in the next one.