Рецепт vLLM с мистралем. Часть 1 / Хабр

На первом этапе нашей заготовки модели мы займемся ее 4-х битным квантованием, а во второй части будем использовать vLLM для запуска парралельных запросов к ней.

Поехали!

Для начала запустим саму Mistral-7b в бесплатном колабе, что сразу вызовет трудности т.к. для базового запуска француженки, бесплатных ресурсов не хватит (по умолчанию необходимо 15GB RAM, а дают всего 12GB RAM). Чтобы обойти данное ограничение будем использовать метод 4-битного квантования, который позволяет значительно уменьшить размер модели без потери её производительности. Разберёмся, как загрузить модель в 4-битном формате, изучим все варианты её конфигурации и запустим её для выполнения инференса.

Важно отметить, что этот метод можно применять к любой модели, которая поддерживает аргумент device_map (то есть загрузку с помощью библиотеки accelerate).

Начнем с установки и импорта библиотек (установим некоторые крайние билды прям с гитхаба)

!pip install -q -U bitsandbytes
!pip install -q -U git+https://github.com/huggingface/transformers.git
!pip install -q -U git+https://github.com/huggingface/peft.git
!pip install -q -U git+https://github.com/huggingface/accelerate.git

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

Определим параметры BitsandBytesConfig трансформеров и выделим важные строки кода:

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True, 
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

Q — Quantization (Квантование) + LoRA — Low-Rank Adaptation (Низкоранговая адаптация) — Q — **Quantization (Квантование)** + **LoRA** — **Low-Rank Adaptation (Низкоранговая адаптация)**

Ссылка на статью QLoRA - https://arxiv.org/abs/2305.14314

bnb_4bit_quant_type="nf4": 4-битное квантование поддерживает два типа: FP4 и NF4. Тип NF4 (Normal Float 4) был представлен в статье о QLoRA.
bnb_4bit_compute_dtype=torch.bfloat16: Этот параметр задает точность вычислений. По умолчанию используется float32, но для ускорения работы зададим bfloat16.

Следующим шагом нам нужно залогиниться в HF и указать, какую именно модель мы хотим использовать (её model ID), и затем загрузить её, применив те параметры квантования, которые мы задали ранее.

from huggingface_hub import login

#логинимся в хагин фейс

login(token="ВАШ ТОКЕН ИЗ HUGGING FACE")
print("Successfully logged in to Hugging Face!")

#грузим модель

model_id = "mistralai/Mistral-7B-Instruct-v0.1"
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)

Если теперь вывести структуру модели на экран, вы увидите, что большинство обычных линейных слоёв (nn.Linear) заменились на 4-битные версии (bnb.nn.Linear4bit),

print(model)

------вывод------

"""
MistralForCausalLM(
  (model): MistralModel(
    (embed_tokens): Embedding(32000, 4096)
    (layers): ModuleList(
      (0-31): 32 x MistralDecoderLayer(
        (self_attn): MistralAttention(
          (q_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
        )
        (mlp): MistralMLP(
          (gate_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
          (up_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
          (down_proj): Linear4bit(in_features=14336, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): MistralRMSNorm((4096,), eps=1e-05)
        (post_attention_layernorm): MistralRMSNorm((4096,), eps=1e-05)
      )
    )
    (norm): MistralRMSNorm((4096,), eps=1e-05)
    (rotary_emb): MistralRotaryEmbedding()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)"""

Это и есть доказательство, что наша конфигурация с load_in_4bit=True сработала, и модель действительно загружена в 4-битном формате для экономии памяти.

Убедимся что модель полностью поместилась на весь GPU:

model.hf_device_map

------вывод-------

{'': 0}

"""
0 = модель полностью загружена в видеопамять одной видеокарты 
(первой видеокарты, а она у нас всего одна)

"""

Теперь протестируем модель (тестировать будем на английском, т.к. с русским языком у мистрали есть проблемы (грамматика и контекст хромают)

PROMPT= """ ### Instruction: Act as a data science expert.
### Question:
Explain to me what is Large Language Model. Assume that I am a 5-year-old child.

### Answer:
"""

encodeds = tokenizer(PROMPT, return_tensors="pt", add_special_tokens=True)
model_inputs = encodeds.to(device)

generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=True, pad_token_id=tokenizer.eos_token_id)
decoded = tokenizer.batch_decode(generated_ids)
print(decoded[0])


-----------------------вывод-------------------------

<s> ### Instruction: Act as a data science expert.
### Question:
Explain to me what is Large Language Model. Assume that I am a 5-year-old child.

### Answer:
A large language model is like a super smart robot that knows a lot about language. It can understand what people are saying and help them find information or even make up stories for them. Like a magician, but with words instead of tricks!</s>

Успех!

Как видим модель достаточно хорошо отрабатывает.

Раз уж модель Mistral-7B запускается, то использование vLLM — это логичный следующий шаг для существенного ускорения её работы, особенно при обработке нескольких запросов одновременно, чем и займемся в следующей главе!

До встречи во второй части)

Ссылка на колаб