Run the 47B Parameter Mixtral 8x7b Language Model on a Free Colab GPU

In the fast-moving world of large language models, Mixtral 8x7b stands out as a groundbreaking achievement. Released in 2023 by AI research lab SpeechBrain, this massive 47 billion parameter model leveraged an innovative Mixture of Experts (MoE) architecture to achieve state-of-the-art performance on a wide range of natural language tasks, dethroning previous open-access champions like GPT-3.5 and Anthropic‘s Claude.

But such power typically comes at a steep computational cost. Running inference on Mixtral 8x7b requires an extraordinary amount of GPU memory and processing power, putting it out of reach for most independent researchers and hobbyists.

Or does it? Thanks to some brilliant work from the open-source community, it‘s now possible to run this behemoth of a model on a humble Colab notebook with a single T4 GPU. In this post, we‘ll dive into how this is achieved and walk through the process step-by-step. By the end, you‘ll be generating mind-blowing text outputs with Mixtral, all without spending a dime.

Understanding Mixtral‘s MoE Architecture

First, let‘s make sure we grasp what makes Mixtral tick under the hood. The key innovation is using a Mixture of Experts (MoE) architecture instead of a standard dense model.

In a dense model like GPT-3, every input token gets processed by every layer of the network in sequence. The model‘s capacity scales with the number and size of these layers. Training and running bigger dense models delivers better performance but incurs skyrocketing computational costs.

MoE takes a different approach, sometimes called "conditional computation." Instead of one big network, an MoE model contains many smaller "expert" networks, each of which specializes in handling a particular type of input. For each input token, a lightweight "router" network decides which experts to activate, and only those experts update their computations. The router then combines the experts‘ outputs to produce the final result for that token.

This allows MoE models to scale parameters much more efficiently than dense models, since not every parameter is used for every input. An MoE model like Mixtral can contain many more parameters than a dense model while still being faster to train and run. The router learns to activate only the most relevant experts for a given input, keeping computation sparse.

Mixtral 8x7b contains 8 experts, each equivalent to the 7 billion parameter dense Mistral model. Altogether, it has a whopping 47 billion parameters, but the MoE architecture allows inference to run using only a fraction of those parameters at a time. This is the key to how we‘ll get it running on Colab‘s free T4 GPU.

Challenges of Running Mixtral 8x7b on Colab

Even with the efficiencies of MoE, Mixtral 8x7b is still a beast of a model. The full uncompressed weights take up over 86GB. Colab‘s free T4 GPU instances only offer 16GB of VRAM.

Trying to load the whole model at once would quickly run out of memory. We need a way to compress the model‘s size and cleverly manage which weights are loaded at any given time.

Techniques to Run Mixtral on Colab

A few key techniques make it possible to run Mixtral on Colab‘s free GPUs:

1. Quantization

The first step is to quantize the model‘s weights. Quantization means decreasing the numerical precision of the weights, using fewer bits to store each number. The full Mixtral model uses 16-bit floating point numbers. By quantizing this down to 2-4 bits, we can shrink the model size by 4-8x with minimal loss in quality.

Specifically, the experts are quantized down to 2-bits (more compression because they are activated sparsely), while the router and other layers use 4-bits (their outputs are crucial for every timestep). This squeezes Mixtral down to about 17.5GB total in exchange for a manageable 5% loss on benchmark scores. Much more feasible!

2. Least Recently Used (LRU) Caching

With the model quantized, the next challenge is fitting it into GPU memory. Even quantized, the model is a tight squeeze for 16GB of VRAM. We can‘t quite fit all 8 experts at once.

The solution is a caching system for experts. At any time, only the 4 most recently used experts are actually loaded on the GPU. As different tokens come in, the router determines which experts they need. If the required expert is already loaded, great! Inference will be fast. If not, the needed expert has to be swapped in, replacing whichever currently loaded expert was used least recently.

This Least Recently Used (LRU) caching policy ensures the most relevant experts are usually ready to go on the GPU. It does introduce some latency when an expert needs to be loaded, but that cost is far outweighed by the memory savings of not keeping all experts loaded at once.

3. Speculative Expert Loading

To further reduce the latency introduced by LRU caching, the Mixtral engineers came up with a nifty trick called speculative expert loading.

The key insight is that, while the router network decides which experts to use for a given token, we can make an educated guess which experts will likely be needed for the next token the network generates. By looking at the output activations from the layers before the router, it‘s possible to "forecast" which experts the router will assign with high accuracy.

Rather than waiting for the router to determine which experts are needed and swapping them in on the fly, the speculative loading system proactively fetches experts it predicts will be needed for the next token and gets them loaded into the GPU cache. If the guess is correct, those experts will be ready to go with no latency. Mispredictions simply fall back to normal LRU caching.

Speculative loading will typically fetch 1-2 experts per layer while the current token is being processed, so they are ready by the time the next token is generated. This optimization approximately doubles the inference speed compared to vanilla LRU caching.

Running Mixtral 8x7b on Colab Step-by-Step

Now that we understand the techniques used to run Mixtral on Colab, let‘s walk through the process hands-on!

Colab‘s free 16GB T4 GPU instances are just powerful enough to run the quantized, cached, speculative-loaded Mixtral 8x7b model. Make sure to select a T4 instance before going further.

1. Clone the colab-mixtral GitHub repo and install dependencies:

!git clone https://github.com/dvmazur/mixtral-offloading.git --quiet
!cd mixtral-offloading && pip install -q -r requirements.txt

Clear out the outputs and restart the runtime to prevent memory leaks.

2. Download the quantized Mixtral 8x7b model:

!huggingface-cli download lavawolfiee/Mixtral-8x7B-Instruct-v0.1-offloading-demo \
  --quiet --local-dir \Mixtral-8x7B-Instruct-v0.1-offloading-demo  

This downloads a quantized version of the Mixtral-Instruct variant. Instruct models are fine-tuned to follow instructions, making them more flexible and controllable than the base Mixtral model.

3. Import dependencies and configure the model:

import sys
sys.path.append("mixtral-offloading")

import torch
from torch.nn import functional as F
from hqq.core.quantize import BaseQuantizeConfig
from huggingface_hub import snapshot_download
from transformers import AutoConfig, AutoTokenizer
from transformers.utils import logging as hf_logging
from src.build_model import OffloadConfig, QuantConfig, build_model

model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1" 
quantized_model_name = "lavawolfiee/Mixtral-8x7B-Instruct-v0.1-offloading-demo"
state_path = "Mixtral-8x7B-Instruct-v0.1-offloading-demo"
config = AutoConfig.from_pretrained(quantized_model_name)

device = torch.device("cuda:0")  
offload_per_layer = 4

num_experts = config.num_local_experts
offload_config = OffloadConfig( 
    main_size=config.num_hidden_layers * (num_experts - offload_per_layer),
    offload_size=config.num_hidden_layers * offload_per_layer,
    buffer_size=4, 
    offload_per_layer=offload_per_layer,
)

attn_config = BaseQuantizeConfig(
    nbits=4,
    group_size=64,
    quant_zero=True,
    quant_scale=True,
)
attn_config["scale_quant_params"]["group_size"] = 256

ffn_config = BaseQuantizeConfig( 
    nbits=2,
    group_size=16,
    quant_zero=True, 
    quant_scale=True,
)

quant_config = QuantConfig(ffn_config=ffn_config, attn_config=attn_config)
model = build_model(
    device=device,
    quant_config=quant_config,
    offload_config=offload_config, 
    state_path=state_path,
)

The key settings here:

• offload_per_layer determines how many experts to cache on the GPU per layer. Set to 4 for a 16GB GPU or 2 for a 12GB one.
• attn_config and ffn_config specify the quantization settings for the attention layers and feed-forward expert layers respectively.
• offload_config determines how many total experts to load on the GPU vs offload to CPU memory.

4. Run inference:

from transformers import TextStreamer

tokenizer = AutoTokenizer.from_pretrained(model_name)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

print("User: ", end="") 
user_input = input()
print("\n")

user_entry = dict(role="user", content=user_input)
input_ids = tokenizer.apply_chat_template([user_entry], return_tensors="pt").to(device)
attention_mask = torch.ones_like(input_ids)

print("Mixtral: ", end="")
result = model.generate(
    input_ids=input_ids,  
    attention_mask=attention_mask,
    streamer=streamer,
    do_sample=True,
    temperature=0.9,
    top_p=0.9,
    max_new_tokens=512,
    pad_token_id=tokenizer.eos_token_id,
    return_dict_in_generate=True,
    output_hidden_states=True,
)
print("\n")

This snippet loads the model and runs a simple chatbot-style inference loop. Enter a prompt, and Mixtral will generate a continuation. The TextStreamer is used to stream the output token-by-token for a more interactive experience.

On a T4 GPU, you can expect Mixtral to generate around 2-3 tokens per second once it gets going. The first few tokens are slower as the cache warms up. While not blazing fast, this is remarkably usable performance for such a huge model running on free hardware! Dense models with comparable performance would be impossible to run in this setup.

Limitations and Future Directions

Fitting Mixtral 8x7b onto a free Colab GPU is an impressive feat of engineering. However, it does come with some significant limitations.

Inference is relatively slow, especially for the first few tokens before the cache is loaded. Fine-tuning the model is impractical due to memory constraints. And the tricks used to shrink the model do result in some loss of capability compared to the full uncompressed version.

Looking forward, there are promising avenues to expand the accessibility of huge MoE models like Mixtral. Ongoing work aims to improve the performance of quantized MoE models, making the quality-compression tradeoff more favorable. Newer GPU architectures with more memory and sparsity-focused cores could allow even larger models to run on a single device.

Distributed inference across multiple devices is another exciting direction. Imagine elastically scaling Mixtral to run on however many GPUs are available, from a single Colab notebook to a massive cluster. The DeepSpeed library is working towards this kind of ultra-scalable MoE inference.

But perhaps the most important development is simply how quickly open-access models are advancing. Mixtral raised the bar significantly from GPT-3.5 just earlier this year. At the current pace, even more powerful models may be coming soon, with even better efficiency through architectural innovations like Mixture of Modalities (MoM). The era of massive open-access models is just getting started, and the techniques discussed here will help put them in more hands than ever before.

Conclusion

Mixtral 8x7b is a prime example of how far open-access language models have come. With 47 billion parameters, it matches or exceeds the best proprietary models on a variety of language tasks, taking full advantage of the MoE architecture.

What‘s even more amazing is that, through a clever combination of compression and caching, it‘s possible to run this model on a free Colab notebook. By quantizing weights, dynamically loading experts into GPU memory, and speculative fetching of likely-needed experts, Mixtral‘s massive scale becomes accessible to anyone.

While there are still limitations to this approach, it represents a huge step forward in the democratization of large language models. As the techniques improve and models become more efficient, we can look forward to a future where anyone can harness the full power of natural language AI without breaking the bank.

The ability to run Mixtral on Colab opens up exciting possibilities for research, experimentation, and application. What will you build with it?