使用Mixtral-offloading在消费级硬件上运行Mixtral-8x7B

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

export LD_LIBRARY_PATH="/usr/lib64-nvidia"
export LIBRARY_PATH="/usr/local/cuda/lib64/stubs"
ldconfig /usr/lib64-nvidia

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

import sys
sys.path.append("mixtral-offloading")
import torch
from hqq.core.quantize import BaseQuantizeConfig
from transformers import AutoConfig, AutoTokenizer
from src.build_model import OffloadConfig, QuantConfig, build_mode

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")
num_experts = config.num_local_experts

offload_per_layer = 3

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,
)

import time
tokenizer = AutoTokenizer.from_pretrained(model_name)
duration = 0.0
total_length = 0
prompt = []
prompt.append("Write the recipe for a chicken curry with coconut milk.")
prompt.append("Translate into French the following sentence: I love bread and cheese!")
prompt.append("Cite 20 famous people.")
prompt.append("Where is the moon right now?")
for i in range(len(prompt)):
    model_inputs = tokenizer(prompt[i], return_tensors="pt").to("cuda:0")
    start_time = time.time()
    with torch.autocast(model.device.type, dtype=torch.float16, enabled=True):
        output = model.generate(**model_inputs, max_length=500)[0]
    duration += float(time.time() - start_time)
    total_length += len(output)
    tok_sec_prompt = round(len(output)/float(time.time() - start_time),3)
    print("Prompt --- %s tokens/second ---" % (tok_sec_prompt))
    print(tokenizer.decode(output, skip_special_tokens=True))
tok_sec = round(total_length/duration,3)
print("Average --- %s tokens/second ---" % (tok_sec))