LARRUZ shows how to run Mixtral 8x7B on free Google Colab by offloading layers to RAM and using mixed quantization. The model uses sparse mixture-of-experts, activating only two expert layers per token. The tutorial walks through setup, tokenizer, and generation with performance benchmarks.
Questions this video answers
Around 45 gigabytes of VRAM.
T4 GPU with 16 GB of VRAM (Colab free tier).
Set offload_per_layer to 4 for Colab T4; increase to 5 if you have only 12 GB VRAM (e.g., 3060).
Approximately 16 GB of VRAM and around 11–16 GB of system RAM (video suggests ~11 GB RAM minimum; 12–16 GB recommended).
Attention layers: 4-bit HQQ quantization with group size 64. Experts: 2-bit HQQ quantization with group size 16 (alternative experts at 3-bit is also available).
Chapters
Welcome to the first video of 2024. In this video I'll show you how to run the Mixtral of experts model from Mistral AI on a free tier of Google Colab. Usually this model needs around 45 gigabytes of VRAM to run, but we're going to run this on a free tier of Google collab with 16 gigabytes of VRAM. This is possible thanks to this amazing paper. Fast inference of min Mixture of experts Language models with offloading mixture of experts are sparse in nature. That means in order to process an input prompt at a given time, it uses only a fraction of model layers. In case of the MOE model from Mistral AI, it uses only two layers to process any given token. Theoretically it means that at any given time you just need those two layers to be on the GPU VRAM and you can offload the rest of the layers to the RAM on the system. And this paper shows that you can actually run an MOE by offloading some of its layers even on a free tier of Google Colab. If you are interested, I'll highly recommend to read this paper. There are some very interesting approaches that they have adopted. Here is a table from the paper that shows what type of inference speed to expect effect if you offload part of the model on different consumer grade GPUs. So for example 3080 has 16 gigabytes of VRAM, 3060 has only 12 gigabytes and the T4 GPU on the free tier of Google Collab also has 16 gigabytes of VRAM. Now for this to work we will need a quantized version of the model rather than the full precision model. Here's the official GitHub repo of the project. Here is a quick description of how the algorithm works. They use mixed quantization with hqq. This is a quantization algorithm. Now when they say mixed quantization that means they apply separate quantization schemes to the attention layers and experts. To fit the model into combined GPU and CPU RAM, you can run this on a relatively lower end GPUs. For example, you probably need at least 12 gigabytes of VRAM. But you also need to have enough memory on the system. So you probably need 12 to 16 gigabytes of RAM in order to run this approach. And they deploy their efficient offloading strategy. So in this case each expert per layer is offloaded separately and only brought back to GPU when needed. So as I said in the beginning, they're simply offloading everything to the system. RAM rather than running or keeping it in the gpu. So it's brought back to the GPU only when we need a specific expert. Okay. Next we are going to quickly look at the official Google Colab they have provided. You can download this notebook on your local machine and run it locally. It doesn't have to be in the Google Colab. So in this case I am running this on a T4 GPU. So here you can see I'm just using T4 GPU from Google Colab that's available on the free tier. Now, there are a couple of things you need to think about when you are running this notebook. So first you need around 16 gigabytes of VRAM and around 11 gigabytes of RAM in order to run this notebook. Depending on your hardware, you can play around with it and you can decide to offload more layers of the model to your ram. But you need higher quantity of RAM for that to work. All of this is controlled by this parameter called offload per layer. So in order to work with Google Colab, you need to set the offload per layer to four. But if you want to offload more layers, you can play around with it. In some cases you'll be able to increase this to let's say five or six if you have less VRAM than this. But in that case you will need more system ram. Now let me walk you through the rest of the notebooks. So first we are importing all the packages that we'll need. We are setting up some environment variables within Google Colab. If you're running this locally, you will still need to do these steps. After that we're cloning the corresponding repo and then we are moving to that specific repo and installing all the required packages. Next we are importing all the different packages that we will need. So we are still using transformers from Hugging Face to both load the tokenizer as well as the model. But there are a couple of other packages that we are going to be using. So for example, offload configuration. This will basically control how many layers we are offloading and also define the quantization configuration in there. Next we need to initialize and set our model parameters. So in this case we're going to be using this specific quantized model. Now if you look at it on Hugging Face, you can see there's mixed quantization being used. So for the attention there it's using 4 bit HQQ format with a group size of 64, but for the experts it's quantized to 2 bits so it's not in 4 bits. And again we're using the HQQ format with a group size Of 16. There is also another version in which they have set the quantization level for the X bits to three bits. So you can play around with this one as well and see which one gives you better performance. So here we set all the configurations based on the quantized version of the model that we loading. After that we set the GPU. In this case there is only a single GPU available. Then offloads per layer is set to 4. This will work with Google Colab. If you have only 12 gigabytes of VRAM, something like 3060, then you can set it to 5. So based on that, we are setting the offload configurations in here. We are also setting the quantization configurations. Now with all those configurations we simply build and load the model. Now how do you run the model? So first we need the tokenizer, the tokenizer based on the original unquantized model. So that's why we were using the model name. If you remember, we had set the model name to the unquantized Mistral Mixtral 87B here. So we're going to be using the tokenizer from there. Then we have, we want to have a streaming text after that. So here's a while loop in which we get an input from the user. Then using the tokenizer, we tokenize our input. And once we process our prompt, the way you generate the response is exactly what you would expect. So we call the generate function pass on our input IDs that is generated through the tokenization process. We also pass our attention mask and stuff. Here's the temperature, top max, new token are set to 512 in this case. So here is an example prompt that the user provided. Write a funny poem about Python, please. And here's the poem that it created. Then the user says, thank you, that's cool. Now write a funny poem about LLMs or large language models. And here is the poem from Mixture, and it actually rhymes pretty good. So it states, there once was a model quite large, you see, made by AI for all humanity. It spoke in tongues from Old English to French, but sometimes its answers would leave you in a lurch. It could write a poem, a story or more, but when asked for advice, it would often be a bore. Okay, not bad at all, right? And then the user goes on to say, summarize your poem about Python in two sentences. And here's a quick summary of that poem. Now keep in mind the speed of generation is not great. It's about 1 1/2 to 22 tokens per second, although processing of the actual prompt takes a while but nevertheless we are able to run a 45 billion parameter model on a 16 gigabytes of VRAM which is pretty crazy to think about. I highly recommend to give it a shot. I'll put all the links in the video description. Just keep in mind the model takes a long time to download so I'm not sure if there's some issue with the hugging face repo on which it's hosted, but it does take quite a while to download the model. I have seen sometimes it timed out on me. The Google collab would restart so you will need to consider those things. So if you have the GPU I'd highly recommend just download the code, run it locally and it seems to be working. I hope you found this video useful. Consider subscribing to the channel and liking the video. Thanks for watching and as always see you in the next one. Thanks.
