Привет, Хабр! Меня зовут Владимир, и это четвёртая часть цикла статей по написанию и обучению небольшой decoder-only LLM с нуля. В третьей части мы собрали LinguaLaboratoriumMechanicus, обучили base-модель на корпусе данных по Warhammer 40K и получили модель, которая умеет продолжать текст. Данная статья целиком посвящена этапу SFT - дообучению на датасете “вопрос - ответ”, чтобы модель могла вести диалог с пользователем, а не просто дописывать за него фразы.

Содержание цикла

1. Подготовка и токенизация данных

2. Трансформер

3. Сборка и обучение LLM

4. SFT этап - дообучение на вопрос-ответ (вы находитесь здесь)

5. Интеграция с Hugging Face

Содержание может меняться и дополняться ссылками по мере написания

Зачем нужен SFT

После pre-train модель уже умеет предсказывать следующий токен. Но если задать ей вопрос “Кто такой Император?”, она не ответит - она продолжит текст так, будто это начало фразы. Модель может выдать что-то вроде “- спросила она. Как он может не быть Богом?” - грамматически верно, но это не диалог.

Pre-train учит модель языку, стилю, даёт запоминать информацию корпуса. SFT этап учит формату поведения, в нашем случае это формат вопрос - ответ. Для этого нужны пары “вопрос - ответ” и другой способ считать loss.

Грубо говоря, если на этапе претрейна мы спрашивали “какой токен идёт дальше?”, то на этапе SFT вопрос уже “какой токен идёт дальше в ответе на этот вопрос?”.

Механизм обучения тот же - cross-entropy на следующий токен. Меняются данные, формат промпта и то, что loss считается только на токенах ответа модели.

Данные для SFT

Для pre-train данные, условно, подойдут любые: чем больше сырого текста, тем лучше. Для SFT объём уже не главное - нужен не корпус абзацев, а примеры желаемого поведения модели. Таких датасетов бывает несколько видов. Классический instruct датасет - это когда модель учат следовать инструкциям, например переводить текст или исправлять грамматику. Есть диалоговый датасет - многоходовые диалоги с чередованием ролей. Также существуют датасеты для обучения моделей агентной работе (в частности, вызову инструментов). Для своей модели я выбрал упрощённый вариант чата - один вопрос и один ответ на него.

Наверное, это был самый трудный этап, во время которого я несколько раз хотел забить на затею с обучением модели. Данные я собирал на тематических форумах, в статьях, когда надоело, то нагенерировал немного синтетики из лорных материалов с помощью нейронки. В итоге собрал несколько JSON файлов формата:

[
  {
    "messages": [
      {"role": "user", "content": "Какие цвета доминируют в геральдике и окраске техники Кровавых Ангелов?"},
      {"role": "assistant", "content": "Доминирующим цветом их брони является вермильоново-красный, а отдельные украшения и детали сияют золотом."}
    ]
  }
]

Каждый элемент списка - один диалог. Внутри - массив messages с ролями user и assistant.

Чат-формат и токены ролей

Чтобы модель понимала, где вопрос, а где ответ, нужен особый текстовый формат. Для этого используют специальные токены-маркеры:

<|user|>
Какие цвета доминируют ...
<|assistant|>
Доминирующим цветом их брони ...

Каждая реплика начинается с токена роли (<|user|> или <|assistant|>), затем перевод строки и текст. Специальные токены <|user|> и <|assistant|> добавляются в словарь как отдельные единицы - модель учится однозначно отличать границы ролей, а не угадывать по обычным словам из корпуса.

Расширение токенизатора

Базовый токенизатор из первой части этих токенов не знает. Добавим их отдельным скриптом:

def extend_tokenizer(base_path: str | Path, out_path: str | Path) -> int:
    out_path = Path(out_path)

    tokenizer = AutoTokenizer.from_pretrained(Path(base_path))
    added = tokenizer.add_tokens(
        ['<|user|>', '<|assistant|>'],
        special_tokens=True)
    out_path.mkdir(parents=True, exist_ok=True)
    tokenizer.save_pretrained(out_path)
    return len(tokenizer)

Словарь вырос на 2 токена, а это значит, что и embedding-слой, и выход модели тоже должны вырасти, но об этом чуть позже. А пока соберём QA датасет

QA Dataset

Опять немного теории о подготовке датасета. В pre-train этапе мы учили модель предсказывать следующий токен, поэтому в pre-train датасете target данные - это смещённые на одну позицию входные токены. На SFT этапе модель также видит весь диалог (иначе не поймёт контекст вопроса), но учится предсказывать должна только токены ответа ассистента. Для этого в PyTorch используется стандартный приём: в таргете на “ненужных” позициях ставим -100, а в cross_entropy передаём ignore_index=-100.

Теперь код. Конструктор отличается от pre-train датасета тем, что массива теперь надо хранить два:

class ChatQADataset(Dataset):
    def __init__(
            self, tokenizer: PreTrainedTokenizerBase, json_path: str | Path, max_length: int = 1024, force_reprocess=False):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.pad_token_id = tokenizer.pad_token_id

        self.processed_dir = Path(json_path) / 'processed'
        self.processed_dir.mkdir(exist_ok=True)
        self.inputs_file = self.processed_dir / 'input_blocks.npy'
        self.targets_file = self.processed_dir / 'target_blocks.npy'
        if force_reprocess or not self.inputs_file.exists():
            self._preprocess_data(json_path)
        self.input_blocks = np.load(self.inputs_file, mmap_mode='r')
        self.target_blocks = np.load(self.targets_file, mmap_mode='r')

Логика сборки последовательности и масок реализована в методе ChatQADataset._build_sequence:

class ChatQADataset(Dataset):
    # Предыдущий код
    def _build_sequence(self, messages: list[dict]) -> tuple[list[int], list[int]]:
        input_ids: list[int] = []
        train_positions: set[int] = set()

        for msg in messages:
            role = msg['role']
            header_ids = self.tokenizer.encode(f'<|{role}|>\n', add_special_tokens=False)
            content_ids = self.tokenizer.encode(msg['content'].strip(), add_special_tokens=False)

            content_start = len(input_ids) + len(header_ids)
            input_ids.extend(header_ids)
            input_ids.extend(content_ids)

            if role == 'assistant':
                for i in range(content_start - 1, content_start + len(content_ids) - 1):
                    train_positions.add(i)

        input_ids.append(self.tokenizer.eos_token_id)
        if messages and messages[-1]['role'] == 'assistant':
            train_positions.add(len(input_ids) - 2)

        labels = [-100] * len(input_ids)
        for i in train_positions:
            if i + 1 < len(input_ids):
                labels[i] = input_ids[i + 1]

        return input_ids, labels

Разберём по шагам.

Сначала для каждого сообщения кодируем заголовок роли (<|user|>\n или <|assistant|>\n) и текст отдельно. Это нужно, чтобы точно знать, где начинается контент ассистента.

Если роль - assistant, добавляем позиции токенов в train_positions. В конце добавляем EOS и формируем labels: везде -100, кроме позиций из train_positions, при этом на i-тую позицию сразу записываем следующий токен input_ids[i + 1] (метод обучения не меняется - предсказание следующего токена).

Собираются диалоги в методе ChatQADataset._preprocess_data:

class ChatQADataset(Dataset):
    # Предыдущий код
    def _preprocess_data(self, json_dir_path):
        json_files = sorted(Path(json_dir_path).glob('*.json'))
        print(f'Найдено {len(json_files)} JSON файлов')

        input_blocks: list[list[int]] = []
        target_blocks: list[list[int]] = []

        for json_path in tqdm(json_files, desc='Обработка файлов'):
            with open(json_path, 'r', encoding='utf-8') as f:
                dialogs = json.load(f)
            for dialog in dialogs:
                input_ids, labels = self._build_sequence(dialog['messages'])
                pad_len = self.max_length - len(input_ids)
                ids = input_ids + [self.pad_token_id] * pad_len
                labs = labels + [-100] * pad_len
                input_blocks.append(ids[:-1])
                target_blocks.append(labs[:-1])
        np.save(self.inputs_file, np.array(input_blocks, dtype=np.int32))
        np.save(self.targets_file, np.array(target_blocks, dtype=np.int32))

Диалоги грузятся по одному, предобрабатываются в buildsequence и дополняются PAD токенами до заданной длины контекста (таргет для PAD токенов заполняем -100). Я пробовал “склеивать” диалоги в одну последовательность через EOS токен (чтобы уменьшить накладные расходы), но чисто субъективно показалось, что результат был хуже (хотя loss вёл себя лучше).

SFT-обучение

Загрузка Base модели и расширение словаря

SFT начинается не с нуля, а с весов модели после pre-train этапа. Только один-в-один базовая модель нам не подойдёт, так как в токенизатор мы ещё два спец токена добавили, поэтому делаем следующее: из модели копируем веса слоя получения эмбеддингов и выходного слоя (те слои, где фигурирует параметр vocab_size). Далее создаём слои новой размерности и переписываем в них веса из старых слоёв, а новые (дополнительные) веса заменяем средним значением (чтобы не совсем случайные коэффициенты были). В коде это выглядит так:

def resize_vocab(model: LinguaLaboratoriumMechanicus, new_vocab_size: int) -> LinguaLaboratoriumMechanicus:
    old_vocab_size = model.vocab_size
    if new_vocab_size == old_vocab_size:
        return model

    emb_dim = model.emb_dim
    old_emb = model.token_emb.weight.data
    old_head = model.out_head.weight.data

    model.token_emb = nn.Embedding(new_vocab_size, emb_dim)
    model.out_head = nn.Linear(emb_dim, new_vocab_size, bias=False)

    model.token_emb.weight.data[:old_vocab_size] = old_emb
    model.token_emb.weight.data[old_vocab_size:] = old_emb.mean(dim=0)
    model.out_head.weight.data[:old_vocab_size] = old_head
    model.out_head.weight.data[old_vocab_size:] = old_head.mean(dim=0)
    model.vocab_size = new_vocab_size

    return model

Цикл обучения

Цикл обучения в целом повторяет третью часть: AdamW, cosine scheduler с warmup, gradient clipping, сохранение чекпоинтов после каждой эпохи. Основные отличия - данные (c маской в loss), другой learning rate и способ оценки результата.

Learning rate на SFT в 30 раз ниже (1e-5 вместо 3e-4 на претрейне), чем на претрейне - модель уже обучена и нам нужно слегка “подкрутить” поведение, а не переучивать с нуля. Слишком большой LR на этом этапе легко “сносит” то, что модель выучила на pre-train.

Метод оценки генерации должен повторять формат обучения: вопрос пользователя + маркер начала ответа ассистента. Метод run_eval в SFT цикле собирает промпт и вызывает generate из третьей части:

@torch.no_grad()
def run_eval(model, tokenizer, cfg: Config) -> None:
    model.eval()
    for question in cfg.eval_questions:
        prompt = format_messages([{'role': 'user', 'content': question}])
        prompt += '\n<|assistant|>\n'
        output = generate(
            model, tokenizer, prompt,
            max_new_tokens=cfg.eval_max_new_tokens,
            temperature=cfg.eval_temperature,
            top_k=cfg.eval_top_k,
            device=cfg.device,
        )
        print(f'Вопрос: {question}')
        print(f'Ответ: {output}')

prompt += '\n<|assistant|>\n' как раз даёт модели понять, что сейчас её очередь генерировать ответ. Без этого она может просто продолжить вопрос или уйти в сторону.

Результаты

SFT этап обучения занял всего 20 минут (около 80 секунд на эпоху).

Loss при этом вёл себя “правильно”: с 3.7 на первой эпохе до 2.1 на пятнадцатой. Уже к середине обучения модель уверенно отвечает в формате ассистента, однако, хоть loss и продолжает падать, генерация начинает деградировать: появляются зацикливания, тавтологии и подмена сущностей.

По сути модель всё лучше подгоняется под тренировочные пары вопрос–ответ, но на маленьком QA-датасете это быстро превращается в переобучение: кросс-энтропия на маске ответа падает, а полезность ответа для читателя - нет. Для такой ситуации разумнее выбирать чекпоинт не по минимальному loss, а по качеству генерации на фиксированных вопросах.

Сравним поведение base-модели и SFT-версии на одном вопросе "Что такое Гибельный шторм?":

БАЗОВАЯ: Я не знаю, - ответил Жиллиман. - Я не знаю, что это такое. - Я не знаю. - Тогда я не знаю. - Нет, - сказал Жиллиман. - Я не знаю. Я не знаю. - Тогда я должен это сделать. Жиллиман положил руку на плечо примарха. - Я не могу. - Тыне можешь.

ДООБУЧЕННАЯ: Из-за этого шторма и его действия на Терре лоялисты стали жертвами, которые не были случайными. Они были отрезаны от других легионов, а их решимость была непоколебимой

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

Формат ответа модель усвоила: она перестала просто дописывать книжный абзац и начала отвечать в роли ассистента. А вот с фактами по лору всё ещё беда - иногда генерируется что-то очень странное или вообще не в тему.

Тем не менее, пайплайн рабочий: base -> SFT -> чат-модель.

Код обучения тут

Base-модель тут

Instruct-версия тут

А в следующей части цикла как раз разберём, как нашу модель на PyTorch упаковать для Hugging Face и загрузить на Hub, чтобы ей можно было пользоваться с помощью замечательного transformers.AutoModelForCausalLM.