LLMClone: как клонировать себя в Telegram

У меня, как и у многих, довольно много чатов в телеграмме. Иногда просто нет времени (а иногда и не хочется) отвечать на некоторые сообщения. Именно так возникла идея создания виртуального клона. В статье рассматривается простая идея, состоящая в том, чтобы зафайнтюнить языковую модель на личных сообщениях, выгруженных из Telegram-чатов. Возможно, в дальнейшем такой клон сможет общаться за вас

В статье показано обучение квантованной модели именно в коде (в других статьях по созданию клонов, например здесь, используются скрипты Lit-GPT, модель обучается в 16 бит) и создана небольшая библиотека, чтобы каждый мог попробовать обучить клона в гугл коллабе (в аналогичной статье используется A100 40GB). Приятного чтения!

Мозг (языковая модель)

В качестве базовой модели или "мозга" для клона была выбрана decoder-only модель Mistral с 7 миллиардами параметрами, которая по производительности выигрывает у 13- миллиардной LLaMA-2.

Одной из архитектурных особенностей Mistral-7B является то, что в ней используется SWA (sliding window attention), в котором каждый слой отслеживает предыдущие 4096 скрытых состояний. Основной причиной, по которой эта технология была использована в Mistral является линейная вычислительная стоимость O(sliding_window.seq_len). На практике (вместе с измененными FlashAttention и xFormers) SWA позволяют увеличить скорость в 2 раза при длине последовательности в 16k и с окном в 4k токенов.

Каждый токен связан с W токенами из предыдущего уровня (здесь W = 3). Токены за пределами скользящего окна влияют на предсказание следующего слова. На каждом слое механизма внимания информация может перемещаться вперед на W токенов. Следовательно, после k слоев внимания информация может продвигаться вперед на k × W токенов.

Подготовка данных

Для начала выгружаем json со всеми личными сообщениями. Для этого нужно зайти в настройки Advanced и нажать export telegram data:

Отмечаем нужное (только Personal chats) и выгружаем сообщения в формате json:

Далее можно переходить к препроцессингу сообщений. Объединяем сообщения с отправителем. Никнеймы всех юзеров (кроме клонируемой личности) заменяем на "User":

import pandas as pd
from datasets import Dataset, load_from_disk

def process_chats(file_path: str) -> List[str]::
    df = pd.read_json(file_path)

    messages = []
    for sample in df["chats"]["list"]:
        for row in sample["messages"]:
            if row["text"] != '':
                username = row['from']
                if username != "Alan":
                    username = "User"
                if username == "Alan":
                    username = "Clone"  
                message = f"{username}: {row['text']}"
                messages.append(message)
    
    return messages

Объединяем несколько сообщений подряд от одного и того же пользователя в одно большое сообщение:

merged_messages = []
current_user = ''

for message in messages:
    if message.startswith('User:'):
        if current_user != 'User':
            current_user = 'User'
            merged_messages.append(message)
        else:
            merged_messages[-1] += '\n' + message[len('User: '):]
    else:
        if current_user != 'Clone':
            current_user = 'Clone'
            merged_messages.append(message)
        else:
            merged_messages[-1] += '\n' + message[len('Clone: '):]

Предыдущие блоки кода можно объединить в одну функцию, но для наглядности оставим как есть. Разбиваем диалоги между User и Clone в группы по 5 сообщений и создаем экземпляр класса Dataset:

size = 5
num_steps = len(merged_messages)/5
samples = ("\n".join(merged_messages[i*size:(i+1)*size]) for i in range(round(num_steps)))

df = pd.DataFrame({"prompt": samples})
dataset = Dataset.from_pandas(df)
dataset.save_to_disk("clon_conversations")

print(dataset)

# >>> Dataset({
# >>>     features: ['prompt'],
# >>>     num_rows: 2460
# >>> })

print(dataset[1602].get("prompt"))

# >>> User: Постараюсь в ближайшее время это уже доделать
# >>> Clone: Давай брат. Как там с тестами
# >>> User: Мне чёто не нравится пока чё происходит.
# >>> Clone: Помнишь пакеты распознавания речи?
# >>> User: Помню

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

Обучение модели в int 4 с QLoRA

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

Тема LoRA много раз затрагивалась на хабре, поэтому в двух словах: к слоям языковой модели прикрепляем адаптеры низкого ранга и обучаем только их. В случае, когда мы хотим обучить квантованную в 4 бит модель, на помощь приходит метод QLoRA.

В сравнении со стандартным файнтюнингом в 16 бит, QLoRA значительно сокращает использование памяти. Например, с помощью этого метода можно запустить обучение Llama-7B в Google Colab. Модель в таком случае будет занимать около 3,5 гигабайт.

QLoRA использует 4-битное квантование для сжатия предобученной языковой модели. Затем параметры языковой модели замораживаются, и в модель добавляется относительно небольшое количество обучаемых параметров в виде адаптеров с низким рангом. Во время тонкой настройки QLoRA переносит градиенты через замороженную 4-битную квантованную языковую модель в адаптеры. Слои LoRA - это единственные параметры, которые обновляются во время обучения.

QLoRA имеет один тип данных хранения (обычно 4-битный NormalFloat) для весов базовой модели и тип данных BrainFloat 16, используемый для выполнения вычислений. QLoRA деквантизирует веса от типа данных хранения до вычислительного типа данных для выполнения прямого и обратного проходов, но градиенты вычисляются только для 16 битных адаптеров. Веса деквантизируются только тогда, когда они необходимы, поэтому использование памяти остается низким во время файнтюнинга и инференса. Информация взята из блога HF

Перейдем к тонкой настройке нашего "клона". Для начала импортируем все необходимое из библиотек transformers и peft:

from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)

from peft import (
    LoraConfig, 
    get_peft_model, 
    prepare_model_for_kbit_training,
    TaskType
    )

Для квантизации модели в 4 бит используется библиотека BitsAndBytes, которая уже интегрирована в transformers. Модель можно квантовать двумя способами:

С использованием аргумента load_in_4bit:

model = AutoModelForCausalLM.from_pretrained(checkpoint, 
                                             load_in_4bit=True,
                                             device_map="auto")

Продвинутое использование (BitsAndBytesConfig):

• bnb_4bit_compute_dtype. Аргумент позволяет менять тип вычислительных данных на bf16 для ускорения. Дефолтное значение равно fp32

• bnb_4bit_quant_type. 4-битная квантизация может производиться с 2 различными типами квантования: FP4 и NF4 (используется по умолчанию). Тип NF4 (NormalFloat 4) и представлен в статье QLoRA. NormalFloat это тип данных, адаптированный для весов, которые были инициализированы с помощью нормального распределения. Подробнее можно прочитать здесь

• bnb_4bit_use_double_quant. Вложенное квантование снижает расход памяти - по эмпирическим наблюдениям, это позволяет зафайнтюнить llama-13b на 16 ГБ NVIDIA-T4 с длиной последовательности 1024, размером батча 1 и шагом накопления градиента (gradient accumulation) 4. Чтобы включить эту функцию, добавляем bnb_4bit_use_double_quant=True при создании конфига

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

Загружаем и квантуем базовую модель:

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

Затем мы должны применить некоторую предварительную обработку к модели, чтобы подготовить ее к обучению в 4 бит. Для этого используем prepare_model_for_kbit_training

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

Прикрепляем адаптеры

peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    target_modules=[
              "q_proj", 
              "k_proj",
              "v_proj",
              "o_proj"],
)

model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

Подгружаем датасет локально и токенизируем:

dataset = load_from_disk("clone_conversations")
dataset = dataset.map(lambda example: tokenizer(example["prompt"], max_length=256), batched=True)
dataset = dataset.train_test_split(0.1, 0.9)

Запускаем обучение с оптимизатором paged_adamw_8bit, который позволяет избежать утечек памяти, которые могут возникнуть при использовании gradient_checkpointing_enable

collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

training_args = TrainingArguments(
    output_dir="llama",
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=4,
    warmup_steps=2,
    logging_steps=100,
    save_steps=1000,
    learning_rate=2e-4,
    optim="paged_adamw_8bit",
    fp16=True,
    num_train_epochs=10,
    ddp_find_unused_parameters=False,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=collator,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"]
)

trainer.train()

model.save_pretrained("clone_peft")

Пример инференса:

from peft import PeftModel

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

model = AutoModelForCausalLM.from_pretrained(
                checkpoint, 
                quantization_config=bnb_config,
                device_map="auto"
            )

model = PeftModel.from_pretrained(model, "clone_peft")
tokenizer = AutoTokenizer.from_pretraiend(checkpoint)

def generate_sample(
            prompt,
            num_return_sequences=1,
            max_new_tokens=128,
            max_length=1024
        ):
        
        input_ids = tokenizer(
            prompt,
            max_length=max_length,
            truncation=True,
            return_tensors="pt"
        ).input_ids

        tokens = model.generate(
            input_ids=input_ids,
            max_new_tokens=max_new_tokens,
            num_beams=num_return_sequences
        )

        decoded_tokens = tokenizer.decode(
            tokens[0],
            skip_special_tokens=True
        )

        return decoded_tokens

generate_sample("Как дела?")
# >>> Пойдет

LLMClone

Возможно, это излишне (HF Transformers и так является высокоуровневой библиотекой), но если кому- то хочется быстро попробовать создать клона из коробки, я написал небольшую либу LLMClone