Оптимизируем дообучение LLM: теория + гайд

В эпоху стремительного развития искусственного интеллекта большие языковые модели (Large Language Models, LLM) становятся неотъемлемой частью множества приложений – от интеллектуальных чат-ботов до систем анализа данных. 

Так уж сложилось, что эффективное применение больших языковых моделей не обходится без тонкой настройки, потому что базовые модели, обученные на обобщенных данных, могут не учитывать уникальные особенности конкретных задач или доменов. Тонкая настройка позволяет адаптировать модель к специфическим требованиям приложения, что улучшает ее точность [1]. 

Согласно исследованию о модели GPT-3 [2], тонкая настройка на специализированных наборах данных значительно повышает эффективность модели в узконаправленных задачах, таких как медицинская диагностика или юридический анализ.

Однако тонкая настройка может потребовать значительных вычислительных ресурсов, особенно в части использования GPU [4]. В этой статье я расскажу, как настроить небольшую LLM (llama-3.2-3b-instruct) так, чтобы снизить требования к оборудованию и ускорить процесс интеграции модели в бизнес-процессы.

В этой статье мы рассмотрим:

1. Что такое fine tuning

2. Основные методы оптимизации файнтюнинга LLM

3. Применение библиотек transforemers и unsloth для файнтюнинга и его оптимизации на примере практического кейса

Пару слов про кейс

Разумеется, дообучение небольшой LLM не было самоцелью. Я решал задачу по созданию чат-бота, который должен был извлекать по запросу пользователя из Wiki студии разработки программных роботов (RPA) необходимые блоки, перечислять параметры и давать им краткое и структурированное описание, а также примеры использования.

Введение в файнтюнинг

Чтобы перейти к файнтюнингу, следует рассмотреть вопрос, нужен ли он LLM сейчас.

В течение жизненного цикла LLM проходит три основных этапа: предобучение, файнтюнинг и алаймент. Предобучение модели происходит на огромных объемах текстовых данных из разнообразных источников: книги, научные статьи, интернет-сайты и других источники. Цель предобучения -- сформировать базовые навыки обработки и понимания языка, включая грамматику, семантику, общие знания и культурные контексты.

После предобучения модель проходит этап файнтюнинга, на котором она адаптируется под более специфичные задачи. В отличие от предобучения, файнтюнинг предполагает обучение на более ограниченном объеме данных. Файнтюнинг можно поделить на два типа: 

1. Instruction fine tuning (инструкционный фантюнинг): это обучение языковой модели на специально подготовленных примерах, содержащих инструкции для выполнения конкретных, но не специализированных задач. Этот метод позволяет модели лучше понимать запросы и адаптироваться к различным сценариям, следуя заданным инструкциям. 

2. Domain-specific fine-tuning (файнтюнинг под домен): этот процесс включает обучение модели на данных, специфичных для конкретной области (например, медицина, юриспруденция, финансы), чтобы повысить ее точность и релевантность в данной тематике [5].

На заключительном этапе (алаймент) внимание может уделяться согласованности модели с человеческими ценностями, правильности “интерпретации” части промпта. Для достижения таких результатов используется техника, известная как RLHF (обучение с подкреплением с обратной связью от человека), когда модель обучается на основе оценки ответов, которые получены от специально обученной модели вознаграждения или от обратной связи со стороны пользователей [5]. 

Для определения необходимости в файнтюнинге модели стоит оценить текущие задачи и ее состояние. Если модель уже демонстрирует хорошее понимание языка, а также успешно справляется со следованием инструкции, дополнительный инструкционный файнтюнинг может не понадобиться. Однако когда задачи требуют специализированных знаний, то дообучение может значительно улучшить полезность ответов в конкретном домене (правда при этом не стоит забывать, что в таком случае может произойти “катастрофическое забывание” старых знаний [6]).

Основные группы методов файнтюнинга

Еще раз повторим важный тезис: файнтюнинг это процесс адаптации имеющейся модели к новой задаче путем дообучения её на небольшом, связанном с задачей, наборе данных. Он позволяет ликвидировать пробел между универсально обученными моделями и уникальными требованиями специфических доменов. Тут отлично подойдет сравнение - швея подгоняет костюм под определенного человека.

В целом, существует несколько основных техник дообучения модели:

1. Дообучение всей модели целиком на конкретную задачу.

Самый ленивый и неэффективный по требованиям к оборудованию способ. Буквально лишь выставив низкий learning rate и поставив малое количество эпох, запускаем обучение на конкретных данных. Недостатки этого подхода очевидны: помимо того, что под каждую отдельную задачу будет необходима целая отдельная модель, придется обучать всю модель целиком, что может потребовать значительных ресурсов.

2. Заморозка некоторых слоев / обучение части весов модели (репараметризация)

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

3. Добавление адаптеров к модели (аддитивный подход)

Этот метод позволяет донастроить модель для новой задачи без изменения ее основных параметров: то есть мы замораживаем базовую модель (энкодер) и к ней добавляем дополнительные слои. Адаптеры – это (как правило) небольшие дополнительные модули, которые добавляются к основной модели. В этом случае обучаются только параметры адаптеров, что существенно снижает требуемые вычислительные ресурсы и объемы данных. Этот метод, в целом, достаточно эффективен, поскольку позволяет использовать одну и ту же исходную модель для различных задач. Однако использование как правило увеличивает размер модели на вес каждого адаптера.

4. Выбор оптимальных параметров (селективный подход)

Селективные методы файнтюнинга, такие как Fish Mask, позволяют адаптировать LLM к новым задачам, изменяя только критически важные параметры. Этот подход включает анализ и маскирование параметров: создается маска, определяющая, какие веса модели нужно обновить, а какие оставить неизменными. Благодаря этому модель обучается быстрее, так как обучение происходит только по релевантным параметрам, что снижает объем вычислений и помогает избежать «катастрофического забывания» предыдущих знаний [7].

Три основных метода файнтюнинга LLM

Prefix-tuning [8]

Prefix tuning является методом, оптимизирующим небольшие непрерывные префиксные векторы в механизме внимания моделей. При использования данного метода входные данные дополняются контекстом, специфичным для конкретной задачи. Этот контекст, называемый префиксом, представляет собой последовательность обучаемых векторов, добавляемую в виде эмбеддингов, которые встраиваются в механизм внимания на уровне ключей (k, key) и значений (v, value) в каждом слое трансформера.

В отличие от полной настройки параметров модели, здесь изменяются только небольшие векторы, что экономит память и повышает эффективность. 

Применение такого подхода актуально для задач с ограниченными данными, таких как суммаризация текста и перевод, поскольку он позволяет добиться высокой производительности при минимальном количестве настраиваемых параметров.

Реуализация префиксного кодировщика на Torch:

import torch


class PrefixEncoder(torch.nn.Module):
    def __init__(self, num_virtual_tokens, num_layers, token_dim):
        super().__init__()
        self.embedding = torch.nn.Embedding(num_virtual_tokens, num_layers * 2 * token_dim)

    def forward(self, prefix: torch.Tensor):
        return self.embedding(prefix)

Реализация в PEFT:

from peft import PrefixEncoder, PrefixTuningConfig

config = PrefixTuningConfig(
    peft_type="PREFIX_TUNING",
    task_type="SEQ_2_SEQ_LM",
    num_virtual_tokens=20,
    token_dim=768,
    num_transformer_submodules=1,
    num_attention_heads=12,
    num_layers=12,
    encoder_hidden_size=768,
    prefix_projection=False,
)
prefix_encoder = PrefixEncoder(config)

Prompt-tuning [9]

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

В методе Prompt Tuning, вместо изменения параметров самой модели, векторизированные представления промпта настраиваются так, чтобы лучше подходить для решения конкретной задачи. 

В отличие от промпт-инженеринга, промпты здесь не пишутся вручную, а оптимизируются с помощью машинного обучения. Этот метод требует больше вычислительных ресурсов и времени, но создаёт более точные и специализированные представления инструкций для модели. 

Основное различие между prefix-тюнингом и prompt-тюнингом заключается в расположении и обработке параметров. В prompt-тюнинге фиксированные обучаемые токены влияют на модель только на уровне входных данных. В prefix-тюнинге обучаемый префикс добавляется к вниманию, что делает его более гибким для обучения.

Реализация промпт-эмбеддингов на Torch:

import torch

class PromptEmbedding(torch.nn.Module):
    def __init__(self, num_virtual_tokens, token_dim, word_embeddings, init_text=None):
        super().__init__()
        self.embedding = torch.nn.Embedding(num_virtual_tokens, token_dim)
        
        if init_text:
            init_token_ids = init_text
            num_text_tokens = len(init_token_ids)
            
            if num_text_tokens > num_virtual_tokens:
                init_token_ids = init_token_ids[:num_virtual_tokens]
            elif num_text_tokens < num_virtual_tokens:
                num_reps = -(-num_virtual_tokens // num_text_tokens)
                init_token_ids = (init_token_ids * num_reps)[:num_virtual_tokens]
            
            init_token_ids = torch.LongTensor(init_token_ids).to(word_embeddings.weight.device)
            word_embedding_weights = word_embeddings(init_token_ids).detach().clone().to(torch.float32)
            self.embedding.weight = torch.nn.Parameter(word_embedding_weights)

    def forward(self, indices):
        return self.embedding(indices)

Реализация в PEFT:

from peft import PromptEmbedding, PromptTuningConfig

config = PromptTuningConfig(
    peft_type="PROMPT_TUNING",
    task_type="SEQ_2_SEQ_LM",
    num_virtual_tokens=20,
    token_dim=768,
    num_transformer_submodules=1,
    num_attention_heads=12,
    num_layers=12,
    prompt_tuning_init="TEXT",
    prompt_tuning_init_text="Суммаризируй представленную информацию",
    tokenizer_name_or_path="t5-base",
)
# t5_model.shared -- эмбеддинги T5 Base
prompt_embedding = PromptEmbedding(config, t5_model.shared)

P-tuning [10]

P-тюнинг и промпт-тюнинг во многом схожи, поскольку оба метода позволяют адаптировать модель к конкретным задачам с помощью "обучаемых промптов". 

Отличие в том, что результатом является не просто эмбеддинг промпта, а полноценный кодировщик. Кодировщик, например, LSTM+MLP, позволяет адаптировать векторные промпты, учитывая контекст задачи, и преобразовать их для подачи на вход модели. Благодаря этому подходу можно более гибко и точно «обучать» промпты для улучшения производительности модели на конкретных задачах, а также учитывать различные аспекты текста, повышая качество вывода без изменения самой модели.

Реализация на torch:

import torch


class PromptEncoder(torch.nn.Module):
    def __init__(self, num_virtual_tokens, token_dim, encoder_hidden_size):
        super().__init__()
        self.token_dim = token_dim
        self.input_size = self.token_dim
        self.output_size = self.token_dim
        self.hidden_size = encoder_hidden_size
        self.total_virtual_tokens = num_virtual_tokens

        self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim)
        self.lstm_head = torch.nn.LSTM(
            input_size=self.input_size,
            hidden_size=self.hidden_size,
            num_layers=2,
            dropout=0.1,
            bidirectional=True,
            batch_first=True,
        )
        self.mlp_head = torch.nn.Sequential(
            torch.nn.Linear(self.hidden_size * 2, self.hidden_size * 2),
            torch.nn.ReLU(),
            torch.nn.Linear(self.hidden_size * 2, self.output_size),
        )

    def forward(self, indices):
        input_embeds = self.embedding(indices)
        lstm_output = self.lstm_head(input_embeds)[0]
        return self.mlp_head(lstm_output)

Реализация в PEFT:

from peft import PromptEncoder, PromptEncoderConfig

config = PromptEncoderConfig(
    peft_type="P_TUNING",
    task_type="SEQ_2_SEQ_LM",
    num_virtual_tokens=20,
    token_dim=768,
    num_transformer_submodules=1,
    num_attention_heads=12,
    num_layers=12,
    encoder_reparameterization_type="MLP",
    encoder_hidden_size=768,
)

prompt_encoder = PromptEncoder(config)

LoRA

Метод LoRA (Low-Rank Adaptation), разработанный еще в 2021 году [11], предполагает, что объем необходимых для решения задачи изменений в модели, можно сократить, используя низкоранговые матрицы. Вместо того чтобы менять всю матрицу весов, метод добавляет две небольшие матрицы, которые сохраняют только важные параметры, при этом основная матрица остается неизменной, а изменения происходят только в этих новых матрицах, которые затем объединяются с основной для получения результата.

Представьте, что вы хотите отремонтировать комнату, но вместо того чтобы перекрашивать все стены, менять пол, делать натяжной потолок, вы просто добавляете несколько постеров и коврик, чтобы изменить атмосферу. Стены и пол остаются такими же, а небольшие изменения создают нужный эффект.

Точно так же LoRA не обновляет всю огромную матрицу параметров модели, а добавляет небольшие изменения с нужными “деталями”, оставляя основную матрицу нетронутой.

Этот пример не без иронии, ведь эффект “ремонта” здесь иллюзорен, как и файн тюнинг, направленный на кардинальное изменение поведения модели через небольшое количество данных, малое количество эпох и количество параметров. 

Однако, в целом, в отличии от предыдущих методов, LORA потенциально позволяет провести почти полное переобучение модели при сокращении количества параметров. 

При использовании метода LoRA в нашем распоряжении имеются следующие гиперпараметры.

Ранг (r)  определяет размерность низкорангового подпространства, которое задается парами матриц 𝐴 и 𝐵. Чем меньше значение ранга, тем меньше параметров обновляется у модели, однако слишком малый ранг может привести к плохой обучаемости модели.

Низкий ранг (например,  𝑟 = 4) подходит для задач, схожих с исходной, так как минимально изменяет представления, сохраняя исходные знания модели; средний ранг (например, 𝑟 = 8 или 𝑟 = 16) является компромиссом, обеспечивая умеренную адаптацию для задач, требующих дополнительных знаний, но не кардинальных изменений; высокий ранг (например, 𝑟 = 32) используется для задач, сильно отличающихся от исходной, когда требуются значительные изменения, что повышает качество адаптации, но требует больше ресурсов и вычислительных затрат.

Alpha (α) – коэффициент масштаба, регулирующий вклад адаптера в общую адаптацию модели. При низких значениях модель будет меньше полагаться на адаптированные параметры, а при высоких – больше. В практике часто используется значение α между 8 и 32, хотя оптимальный выбор зависит от объема данных. 

LoRA dropout – это гиперпараметр, который регулирует вероятность случайного отключения нейронов в матрицах 𝐴 и 𝐵, добавленных с помощью LoRA. По аналогии с обычным dropout, он помогает улучшить обобщающую способность модели и уменьшить переобучение, особенно когда адаптация выполняется на небольших или сильно ограниченных по объему данных.

Целевые модули – это слои модели, в которые будут интегрированы дополнительные параметры адаптации (матрицы 𝐴 и 𝐵).

Пример реализации LoRA на torch:

import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, rank=4, alpha=1.0, bias=True):
        super(LoRALinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        self.bias = nn.Parameter(torch.empty(out_features)) if bias else None
        self.lora_A = nn.Parameter(torch.empty(rank, in_features))
        self.lora_B = nn.Parameter(torch.empty(out_features, rank))
        self.weight.requires_grad = False
        if self.bias is not None:
            self.bias.requires_grad = False
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=torch.sqrt(torch.tensor(5.0)))
        if self.bias is not None:
            fan_in = self.in_features
            bound = 1 / torch.sqrt(torch.tensor(fan_in, dtype=torch.float32))
            nn.init.uniform_(self.bias, -bound, bound)
        nn.init.zeros_(self.lora_A)
        nn.init.zeros_(self.lora_B)

    def forward(self, x):
        result = x @ self.weight.T
        if self.bias is not None:
            result += self.bias
        lora_update = (x @ self.lora_A.T) @ self.lora_B.T * self.scaling
        return result + lora_update

Реализация в PEFT:

from transformers import AutoModelForSeq2SeqLM
from peft import LoraModel, LoraConfig

config = LoraConfig(
    task_type="SEQ_2_SEQ_LM",
    r=8,
    lora_alpha=32,
    target_modules=["q", "v"],
    lora_dropout=0.01,
)

model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
lora_model = LoraModel(model, config, "default")

Почему мы всегда выбираем LORA?

Для задачи, требующей эффективного дообучения больших моделей, LoRA представляется собой наиболее подходящий вариант. Этот подход позволяет дообучить большие модели без изменения числа параметров и незначительным влиянием на производительность. При этом сохраняется высокая степень гибкости, так как мы можем прямо указать требуемый объем изменений.

В таблице перечислены основные характеристики каждого варианта, а также их преимущества и недостатки. 

Решение кейса

В качестве модели была выбрана компактная Llama 3.2-3b, что было обусловлено, во-первых, тем, что модели с меньшим количеством параметров, как правило, достаточно плохо обучаются и имеют сильно ограниченное количество “знаний”, а также требованиями к потребляемым моделью ресурсам.

0. Ищем датасет для Instruction tuning

Так как модель Llama 3.2-3b достаточно плохо справляется с обработкой текста на русском языке, допускает опечатки и другие ошибки, нам необходимо инструкционное обучение, а значит и датасет с большим количеством инструкций. К счастью, они уже есть -- как переводные, так и синтетически сгенерированные. И один из самых больших и удачных (по моему опыту датасетов) был сделан командой Vikhr -- это датасет Vikhrmodels/GrandMaster-PRO-MAX.

1. Генерируем датасет 

На этом этапе можно использовать GPT, Huggingface Inference или свой сервер с хорошей большой языковой моделью (для такой сгодится, например, Vikhrmodels/Vikhr-Nemo-12B-Instruct-R-21-09-24).

Код для генерации датасета через HuggingFace:

from huggingface_hub import InferenceClient

client = InferenceClient() 

response = client.chat_completion(
    model="mistralai/Mistral-Nemo-Instruct-2407t",     
    messages=[

        {"role": "system", "content": "<системная инструкцият>"},
        {"role": "user", "content": "<пользовательская инструкция>"}
    ],
    temperature=0.2,
    max_tokens=2048,
    stream=False,
)

Для генерации датасета через HF можно использовать разные модели, но наиболее быстро и удачно можно получить результаты от Mistral Nemo.

Код для генерации датасета через OpenAI API (vllm/llama cpp server):

import os

from openai import OpenAI

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"),)

chat_completion = client.chat.completions.create(
    messages=[
         {
            "role": "user",
            "content": "<системная инструкция>",
        },
        {
            "role": "user",
            "content": "<пользовательская инструкция>",
        }
    ],
    model="gpt-4",

)

Генерация при этом должна проходить в два основных этапа: на первом этапе генерируется основа для инструкций (например, определенные слова или фразы, которые должны быть вставлены в промпт или сами инструкции в форме завершенных промптов).

Также стоит отметить, что не стоит запрашивать от модели слишком большие батчи инструкций или инструкции в формате JSON. Модели в обоих этих случаях генерируют некачественные данные. Намного лучше будет, если попросить модель порассуждать над задачей, а уже после предоставить ответ. Точно также можно порекомендовать избегать большого количества капса в промптах, так как это зачастую заставляет вести модель неадекватно, и порекомендовать использовать максимально конкретные слова и фразы, а не как в примере ниже

Неправильный вариант:

Правильный вариант:

Добавляем рассуждение:

2. Чистим данные

Стоит понимать, что те данные, которые мы сгенерировали, далеки от идеала. Можно использовать разные варианты чистки данных, но мной зачастую используется следующий пайплайн: очистка данных с помощью модели, которая обучена отличать плохие генерации от хороших, а затем удаление дубликатов и фильтрация датасета с помощью метода разнообразия. 

Как сделать такую модель? В самом простом варианте, необходимо сгенерировать любой датасет через хорошую и плохую модель. Однако в расширенном, стоит сказать о том, что и плохие и хорошие модели в зависимости от инструкции и настроек генерации могут допускать одинаковые, в сущности, ошибки: плохие -- чуть больше, хорошие -- чуть меньше, и так или иначе придется пройтись по датасету и провалидировать, после чего можно обучить модель классификации.

Обучение модели классификации:

import torch

import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset
from sklearn.metrics import accuracy_score, f1_score, recall_score

def compute_metrics(p):
    predictions, labels = p
    predictions = predictions.argmax(axis=-1)
    accuracy = accuracy_score(labels, predictions)
    f1 = f1_score(labels, predictions)
    recall = recall_score(labels, predictions)
    return {"f1": f1, 'accuracy': accuracy, 'recall': recall}


tokenizer = AutoTokenizer.from_pretrained("belyakoff/deberta_v2_nli", device='cuda:0')
model = AutoModelForSequenceClassification.from_pretrained("belyakoff/deberta_v2_nli", num_labels=2)

for i in model.deberta.parameters():
    i.requires_grad = False


def preprocess_data(examples):
    return tokenizer(examples['content'], padding='max_length', truncation=True)

df = pd.read_excel('train_val_test.xlsx').dropna()
dataset = Dataset.from_pandas(df[['content', 'label']])
tokenized_dataset = dataset.map(preprocess_data, batched=True)

train_testval_split = tokenized_dataset.train_test_split(test_size=0.5, seed=42)
test_val_split = train_testval_split['test'].train_test_split(test_size=0.5, seed=42)
train_dataset = train_testval_split['train']
val_dataset = test_val_split['train']
test_dataset = test_val_split['test']



training_args = TrainingArguments(
    output_dir="./model-steps",
    eval_strategy="steps",
    learning_rate=3e-5,
    eval_steps=250,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=0.001,
    weight_decay=0.001,
    load_best_model_at_end=True,
    lr_scheduler_type='linear',
    optim='adamw_torch'
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    compute_metrics=compute_metrics
)

trainer.train()
trainer.save_model("./model-final")
print(trainer.evaluate(test_dataset))

Выглядит это, в целом, не страшно, однако вы можете сэкономить время на чистке данных, если возьмете готовый классификатор --  в частности belyakoff/deberta_cls_llama_3.1_70b_llama_3_8b, который обучался с разморозкой четвертого и пятого слоев претренированной модели deepvk/deberta-v1-distill.

Код для модели классификации (скоринга)

from torch import nn, Tensor, no_grad
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSequenceClassification

class Scorer(nn.Module):
    def __init__(
            self,
            model_name: str,
            device: str = 'cuda:0',
            tokenizer_name: Optional[str] = None,
            token: Optional[str] = None,
            *args,
            **kwargs,

    ):
        super().__init__(*args, **kwargs)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name, *args, **kwargs, token=token).to(device)
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name or model_name, token=token)
        self.softmax = nn.Softmax().to(device)

    def predict(self, texts: List[str]):
        tokenized: Dict[str, Tensor] = self.tokenizer(texts, return_tensors='pt', padding=True, truncation=True)
        tokenized['input_ids'] = tokenized['input_ids'].to(self.model.device)
        tokenized['attention_mask'] = tokenized['attention_mask'].to(self.model.device)
        self.model.eval()
        with no_grad():
            logits = self.model(**tokenized).logits
            probs = self.softmax(logits)
            return lambda x: x[-1], probs.cpu().numpy()

Теперь перейдем к тому, как его использовать:

1. Проходимся классификатором по всему датасету и рассчитываем скор “низкого качества” -- вероятность плохой модели.

2. Ранжируем сэмплы по скору, определяем границу “приемлемых” данных.

3. Стираем неправильные сэмплы (ставим пометку на удаление).

Похожим путем применяем метод разнообразия. Однако данный метод проблематично использовать, когда у нас слишком много дубликатов в датасете. Для их удаления можно использовать TF-IDF и косинусную близость.

def remove_duplicates_tfidf(texts: List[str], threshold: float = 0.7) -> List[str]:
    vectorizer = TfidfVectorizer(analyzer='char', ngram_range=(2, 3)).fit_transform(texts)
    similarity_matrix = cosine_similarity(vectorizer)
    unique_texts = []
    for i, row in tqdm(enumerate(similarity_matrix)):
        if all(row[j] < threshold for j in range(i)):
            unique_texts.append(texts[i])
    return unique_texts

Теперь ищем центроиды. В примере ниже приведен код для реализации поиска центроидов через K-Means и TF-IDF (самый простой вариант). Это нужно, чтобы отсечь все слишком внешне похожие тексты.

from sklearn.cluster import KMeans
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

def centroids_al(
        texts: List[str],
        n_clusters: int = 600,
        top_n: int = 20,
        tf_idf_kwargs: Dict = None,
) -> List[str]:
    tf_idf_kwargs = tf_idf_kwargs or {}
    vectors = TfidfVectorizer(**tf_idf_kwargs).fit_transform(texts)
    print(vectors.shape)
    model = KMeans(n_clusters=n_clusters).fit(vectors)
    centroids = model.cluster_centers_
    closest_to_centroid: List[str] = []
    for i in range(n_clusters):
        distances = cosine_similarity(vectors, centroids[i].reshape(1, -1)).flatten()
        closest_to_centroid.extend([texts[idx] for idx in distances.argsort()[-top_n:]])
    return closest_to_centroid

Теперь сделаем более тяжелую чистку через torch и контекстуальные эмбеддинги.

from typing import List

import torch
from transformers import AutoTokenizer, AutoModel
import torch.nn.functional as F

model_name = "belyakoff/deberta_v2_nli"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
model.eval()


def get_embeddings(texts: List[str]) -> torch.Tensor:
    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)
    return embeddings


def cosine_similarity_matrix(embeddings: torch.Tensor) -> torch.Tensor:
    norm_embeddings = embeddings / embeddings.norm(dim=1, keepdim=True)
    similarity_matrix = torch.mm(norm_embeddings, norm_embeddings.T)
    return similarity_matrix


def hierarchical_clustering(embeddings: torch.Tensor, n_clusters: int) -> List[List[int]]:
    clusters = [[i] for i in range(len(embeddings))]
    similarity_matrix = cosine_similarity_matrix(embeddings)

    while len(clusters) > n_clusters:
        max_sim = float('-inf')
        closest_pair = (0, 1)
        for i in range(len(clusters)):
            for j in range(i + 1, len(clusters)):
                sim = sum(similarity_matrix[a, b] for a in clusters[i] for b in clusters[j]) / (
                            len(clusters[i]) * len(clusters[j]))
                if sim > max_sim:
                    max_sim = sim
                    closest_pair = (i, j)
        clusters[closest_pair[0]].extend(clusters[closest_pair[1]])
        del clusters[closest_pair[1]]

    return clusters


def centroids_al(
        texts: List[str],
        n_clusters: int = 600,
        top_n: int = 4
) -> List[str]:
    embeddings = get_embeddings(texts)
    clusters = hierarchical_clustering(embeddings, n_clusters)

    closest_to_centroid: List[str] = []
    for cluster in clusters:
        cluster_embeddings = embeddings[cluster]
        centroid = cluster_embeddings.mean(dim=0)

        distances = F.cosine_similarity(centroid.unsqueeze(0), cluster_embeddings)
        closest_texts_indices = distances.argsort(descending=True)[:top_n]
        closest_to_centroid.extend([texts[cluster[idx]] for idx in closest_texts_indices])

    return closest_to_centroid

После того, как у нас очищены сэмплы или очищены инструкции, собираем итоговый датасет, который должен включать: системный промпт, пользовательский промпт и ответ модели. Важно также не забывать, что модель лучше всего обучать на данных с длиной контекста, который соответствует самой модели.

Примерный результат -- belyakoff/constuct-v1-chat.

3. Обучаем модель 

Для обучения модели, как я говорил ранее, мы будем использовать библиотеку Unsloth, которая позволяет сократить объем потребляемых ресурсов в несколько раз.

Однако данная библиотека, как и ряд других около (да пусть меня закидают помидорами) transformers, имеет одну большую проблему -- плохая совместимость пакетов. Для того, чтобы можно было без проблем провести установку, я оставил здесь requirements.

Сам код для обучения выглядит следующим образом:

import torch

from trl import SFTTrainer
from transformers import TrainingArguments, AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset, concatenate_datasets
from unsloth import FastLanguageModel

max_seq_length = 16000

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="belyakoff/llama-3.2-3b-instruct-fine-tunedt",
    max_seq_length=max_seq_length,
    dtype=None,
    load_in_4bit=True,
)
tokenizer = AutoTokenizer.from_pretrained('belyakoff/llama-3.2-3b-instruct')
tokenizer.padding_side = 'left'
print(tokenizer.pad_token )
tokenizer.add_special_tokens(
  {'pad_token': '<|reserved_special_token_0|>'}
)
print(tokenizer.pad_token)

dataset_1 = load_dataset('Vikhrmodels/GrandMaster-PRO-MAX', split='train').map(lambda x: {"text": tokenizer.apply_chat_template(x['conversation'], tokenize=False)}, batched=True)
dataset_2 = load_dataset('Vikhrmodels/Grounded-RAG-RU-v2', split='train').map(lambda x: {"text": tokenizer.apply_chat_template(x['conversation'], tokenize=False)}, batched=True)

dataset = concatenate_datasets([dataset_1, dataset_2]).shuffle()
example = tokenizer.tokenize(dataset['text'][0], )
print(dataset['text'][0])
print(len(example))

model = FastLanguageModel.get_peft_model(
    model,
    r=256,
    target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
    ],
    lora_alpha=256,
    lora_dropout=0.05,
    bias="none",
    use_rslora=True,
    use_gradient_checkpointing=True,
    max_seq_length=max_seq_length,
)


trainer = SFTTrainer(
    model = model,
    train_dataset = dataset,
    dataset_text_field = "text",
    tokenizer = tokenizer,
    packing = True,
    max_seq_length=max_seq_length,
    args = TrainingArguments(
      learning_rate=2e-5,
      output_dir="./model-steps",
      per_device_train_batch_size=1,
      per_device_eval_batch_size=1,
      gradient_accumulation_steps=8,
      warmup_steps=20,
      save_strategy='steps',
      save_steps=300,
      num_train_epochs=20,
      save_total_limit=6,
      fp16=not torch.cuda.is_bf16_supported(),
      bf16=torch.cuda.is_bf16_supported(),
      logging_steps=1,
      optim="paged_adamw_32bit",
      seed=42,
  ),
)


trainer.train()
trainer.save_model('model-final')

В этом коде нужно отметить одну деталь. В случае если модель не подвергалась файнтюнингу ранее и у нее нет padding_token, необходимо его добавить, чтобы SFTTrainer автоматически не установил eos_token=pad_token: если это случится, модель научится генерировать бесконечные последовательности (а это не очень хорошо).

Кроме того необходимо поменять padding_side на left, поскольку это снижает влияние токенов заполнения на модель, позволяя ей сосредоточиться на актуальных данных в начале последовательности, что особенно важно для LLM и других авторегрессионных моделей, обрабатывающих текст слева направо. 

Обучение в случае llama-3.2 проходило в 2 этапа: обучение на общем датасете и обучение на датасете, который был сгенерирован под задачу. На первом этапе параметры были 1:1 как в примере, во втором -- повышен lora dropout до 0.1, сокращен lora rank и alpha, и количество эпох обучения было сокращено до 1, что объясняется небольшим размером данных.

Таким образом была создана модель belyakoff/llama-3.2-3b-instruct-fine-tuned.

Сам пайплайн ассистента был реализован через RAG с использованием Milvus и эмбеддингов  intfloat/multilingual-e5-large-instruct, рассказ про то, как он работает возможен в следующей статье. Если вас заинтересовала эта тема или остались вопросы, пишите об этом в комментариях, я постараюсь ответить

ССЫЛКИ

1 On the Opportunities and Risks of Foundation Models https://arxiv.org/abs/2108.07258

2 Language Models are Few-Shot Learners https://arxiv.org/abs/2005.14165

3 Scaling Laws for Neural Language Models https://arxiv.org/abs/2001.08361

4 Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism https://arxiv.org/abs/1909.08053

5 Training language models to follow instructions with human feedback https://arxiv.org/abs/2203.02155

6 An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning https://arxiv.org/pdf/2308.08747v3

7 Targeted Efficient Fine-tuning: Optimizing Parameter Updates with Data-Driven Sample Selection https://arxiv.org/abs/2403.08484

8 Prefix-Tuning: Optimizing Continuous Prompts for Generation https://arxiv.org/abs/2101.00190

9 The Power of Scale for Parameter-Efficient Prompt Tuning https://arxiv.org/abs/2104.08691

10 GPT Understands, Too https://arxiv.org/abs/2103.10385

11 LoRA: Low-Rank Adaptation of Large Language Models https://arxiv.org/abs/2106.09685