Добрый день, уважаемые читатели и авторы Хабра!

Сегодня я рад представить вам подробное руководство по обучению модели ruGPT-3.5 13B с использованием датасетов модели Saiga-2/GigaSaiga, технологии Peft/LoRA и технологии GGML. Эта статья призвана стать полезным и практичным ресурсом для всех, кто интересуется машинным обучением, искусственным интеллектом и глубоким обучением, а также для тех, кто стремится глубже понять и освоить процесс обучения одной из самых мощных и перспективных русскоязычных моделей.

В данной публикации мы разберем каждый этап обучения модели, начиная от подготовки данных и заканчивая конвертацией в формат GGML. Буду рад, если мой опыт и знания помогут вам в вашем исследовании и экспериментах в этой захватывающей области!

Мотивация

В свете недавних успехов и инноваций в области больших языковых моделей (LLM), особое внимание уделяется созданию более мощных и эффективных систем, способных обрабатывать и анализировать текст на естественном языке.

Одним из ярких примеров таких систем является модель ruGPT-3.5 на 13 миллиардов параметров, созданная командой Sber AI и опубликованная 20го июля 2023 года. Эта модель представляет собой языковую нейронную сеть оптимизированную под работу с русским языком, она может совершать различные задачи обработки текста, от генерации текста до ответов на вопросы и многого другого. Однако, может не значит умеет, оригинальная ruGPT-3.5 "из коробки" хорошо умеет только продолжать текст предложенный пользователем.

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

На Хабре имеется две хорошие публикации, посвященные данной теме:

• Реально Бесконечное (лето) RuGPT3.5: Генерация новеллы на ходу нейросетью - благодаря этой публикации у меня появилась мысль попробовать выполнить дообучение модели самостоятельно, но мне не нравилось то, что автор использовал квантизированную до 4 бит базовую модель, хотелось чтобы исходник был оригинальным (пусть и в режиме load_in_8bit). Ну и в целом, ruGPT-3.5 способна на нечто большее чем писать новеллы.

• Как (быстро) сделать русский локальный ChatGPT - из этой публикации я узнал о проекте rulm, моделях Saiga (основанной на LLaMA) и GigaSaiga (основанной на ruGPT-3.5), попробовал их и был сильно впечатлён. Но ни эта публикация, ни исходные коды или документация, ни какой бы то ни было источник не показывал, как дообучить именно GigaSaiga пошагово.

А все прочие публикации которые мне попадались на глаза либо ссылались на упомянутые выше, либо были вида "Сбер явил миру ruGPT-3.5". И у меня сложилось впечатление, что ML сообществу более интересна тема дообучения моделей семейства LLaMA, хотя на мой скромный взгляд (и опираясь на опыт личного использования) ламы несколько хуже приспособлены для работы с русским языком.

Ну и в целом, зачем заниматься дообучением LLaMA если про это уже написано десятки если не сотни публикаций? В этом нет ни вызова, ни новизны.

Знакомство с ruGPT-3.5

Началось моё знакомство с данной моделью неспешно, после того как появились первые новости о новой модели от Сбера. На тот момент у меня уже имелась RTX 4090 на 24Гб VRAM от Гигабайт, но даже её памяти не хватало для запуска модели. Поэтому я стал искать различные способы её хоть как-то уместить в память карточки, по ходу дела узнал про библиотеку bitsandbytes, которая расширяет функционал библиотеки transformers, добавляя такие замечательные опции как load_in_8bit и load_in_4bit. Упомянутые опции позволяют выполнять квантизацию "на лету", точнее квантизация происходит в момент загрузки модели в оперативную память видеокарты.

По мотивам указанных изысканий я опубликовал на Дзене в своём блоге пост под названием ИИ в каждый дом! Инструкция по запуску ruGPT-3.5 13B. Мои эксперименты показали, что в режиме 8bit качество её работы в целом приемлемое, генерация текста получается не самая плохая и оперативной памяти карточки хватает.

После этого я понял, что надо уже пробовать выполнить дообучение.

Подготавливаем окружение

Как я упомянул ранее, на Хабре мелькала публикация о проекте rulm, автор данной публикации подробно рассказал о том, как ему удалось собрать большой русскоязычный датасет и выполнить дообучение множества различных моделей, включая LLaMA (2) и ruGPT-3.5.

На GitHub страничке указанного проекта я обнаружил Jupyter Notebook tune_llama_7b.ipynb с подробной инструкцией о том, как дообучить LLaMA 2 7B, но ничего похожего про ruGPT-3.5 не было, хотя и упоминалась модель GigaSaiga, конфигурацию gigasaiga_13b.json которой я решил использовать в качестве основы для своих экспериментов.

И так, создадим директорию в которой будем выполнять все действия:

mkdir ruGPT-3.5-training
cd ruGPT-3.5-training

Для дальнейшей работы нам понадобится Python 3.10, хотя возможно и на 3.11 всё будет отлично работать, не проверял. Помимо этого необходим модуль Python VirtualEnv (предпочитаю это решение, так как not a big fan of Conda) и само собой драйвера Nvidia, включая CUDA (я проводил обучение на 12.2).

Создадим виртуальное окружение и переключимся на него:

python3 -m venv venv
source venv/bin/activate

Установим зависимости, вот пример файла requirements.txt

pip install -r requirements.txt

Теперь клонируем репозиторий rulm:

git clone https://github.com/IlyaGusev/rulm.git

Далее скопируем некоторые конфигурационные файлы:

mkdir {configs,internal_prompts,output,output_ggml}
cp rulm/self_instruct/configs/gigasaiga_13b.json configs/rugpt35_13b.json
cp self_instruct/internal_prompts/gigasaiga.json internal_prompts/rugpt35.json

После чего в файле configs/rugpt35_13b.json подправим поле model_name на ai-forever/ruGPT-3.5-13B.

Тренируем модель

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

Шаг 1 - Подготовка датасетов

Для обучения большинства моделей необходимо иметь датасеты: тренировочный и валидационный, но чтобы их сделать необходимо для начала подготовить какой-то общий, большой датасет, но где взять данные? На помощь опять приходит проект rulm, там имеется три скрипта для создания датасетов из данных заранее подготовленных командой rulm.

• create_chat_set.py

• create_instruct_set.py

• create_short_chat_set.py

Каждый из них собирает специфический вариант датасета используемого для обучения разных версий моделей семейства Saiga, но лично меня больше всего заинтересовал create_chat_set.py, так как он предполагал слияние сразу 7 различных датасетов и готовил их таким образом, чтобы после обучения на них получалась модель типа ChatBot, вот полный список:

• ru_turbo_alpaca

• ru_turbo_alpaca_evol_instruct

• ru_turbo_saiga

• ru_sharegpt_cleaned

• oasst1_ru_main_branch

• gpt_roleplay_realm

• ru_instruct_gpt4

К слову сказать, оригинальная GigaSaiga была обучена на 6 из них, не был задействован датасет gpt_roleplay_realm, в нём обыгрываются забавные и нестандартные игровые сценарии общения модели с пользователем.

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

import subprocess
from pathlib import Path

# Set up paths
content_dir = Path('.').resolve()
train_full_path = content_dir / 'train_full.jsonl'
val_full_path = content_dir / 'val_full.jsonl'

# Run create_chat_set script from rulm
module_directory = Path('rulm/self_instruct').resolve()
subprocess.run(
    ['python', '-m', 'src.data_processing.create_chat_set', str(train_full_path), str(val_full_path)],
    cwd=module_directory,
    check=True
)

# Check if train_full.jsonl exists
if not train_full_path.exists():
    raise FileNotFoundError(f"{train_full_path} does not exist")

Исходный код тут.

Запустим и подождём некоторое время:

python3 1_dataset.py

По итогу в корне проекта появится два файла:

• train_full.jsonl - полный тренировочный датасет, содержит примерно 59 тысяч различных документов оформленных в виде чатов между пользователем и нейросетью.

• val_full.jsonl - полный валидационный датасет, содержит почти 3 тысячи документов.

Шаг 2 - Тренировка модели

Датасеты готовы, а это значит, что можно приступать к обучению самой модели, для этого будут использованы библиотеки torch, transformers и peft. Если в двух словах на этом этапе нам потребуется скачать модель ruGPT-3.5-13B из репозитория на HuggingFace, после чего скопировать конфигурации в папку output, а затем внести в них небольшие правки.

Полный код скрипта 2_train.py приводить не буду, можете изучить его тут, остановимся лишь на этапе непосредственного запуска процесса обучения:

import json
from huggingface_hub import snapshot_download
from pathlib import Path
import subprocess

content_dir = Path('.').resolve()
original_config_path = content_dir / 'configs/rugpt35_13b.json'
model_dir = content_dir / "ruGPT-3.5-13B"
base_model = "ai-forever/ruGPT-3.5-13B"
output_dir = content_dir / 'output'
config_path = content_dir / 'configs/rugpt35_13b_colab.json'

# Paths to datasets
train_full_path = content_dir / 'train_full.jsonl'
train_small_path = content_dir / 'train.jsonl'
train_path = train_full_path  # change to train_full_path if you need
val_full_path = content_dir / 'val_full.jsonl'
val_small_path = content_dir / 'val.jsonl'
val_path = val_full_path  # change to val_full_path if you need

# Download binaries
snapshot_download(repo_id=base_model, local_dir=model_dir, ignore_patterns=["LICENSE", "README.md", ".gitattributes"])

...

# Load configurations
with original_config_path.open('r') as fp:
    config = json.load(fp)

# Colab adjustments
config['trainer']['per_device_train_batch_size'] = 2
config['trainer']['per_device_eval_batch_size'] = 1
config['trainer']['gradient_accumulation_steps'] = 128
config['trainer']['eval_steps'] = 50
config['trainer']['save_steps'] = 50
config['max_tokens_count'] = 1000
#config['model_name'] = str(model_dir)
config['templates_path'] = str(content_dir / 'internal_prompts/rugpt35.json')
config['load_in_8bit'] = True
config['load_in_4bit'] = False

# Demo adjustments
config['trainer']['eval_steps'] = 2
config['trainer']['logging_steps'] = 1
config['trainer']['num_train_epochs'] = 1

with config_path.open('w') as fp:
    json.dump(config, fp, indent=4)

# Run training
module_directory = Path('rulm/self_instruct').resolve()
subprocess.run(
    [
        'python', '-m', 'src.train',
        '--config-file', config_path,
        '--train-file', train_path,
        '--val-file', val_path,
        '--output-dir', output_dir,
        '--report-to', 'none'
    ],
    cwd=module_directory,
    check=True
)

...

По коду видно, что происходит запуск модуля src.train в контексте rulm/self_instruct, на вход передаются опции устанавливающие значения до файлов конфигураций, датасетов и директории в которой будет сложен результат.

Запустим его командой:

python3 2_train.py

На моей RTX 4090 обучение заняло примерно 26 с небольшим часов и потребовало примерно 19Гб VRAM, так что пришлось позакрывать многие приложения использующие видеокарту. Кстати, можно немного уменьшить объём необходимой VRAM до 13Гб, для этого потребуется переключить квантизацию "на лету" в режим load_in_4bit:

config['load_in_8bit'] = False
config['load_in_4bit'] = True

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

В результате работы скрипта в директории output появятся следующие файлы:

• adapter_config.json

• adapter_model.bin

• added_tokens.json

• generation_config.json

• merges.txt

• README.md

• special_tokens_map.json

• tokenizer_config.json

• vocab.json

Нас прежде всего интересуют первые два из списка, в файле adapter_model.bin находятся веса LoRA слоя, а в adapter_config.json, конфигурация, которая содержит информацию о том для какой модели создан указанный LoRA слой, как его применять, на какие веса оригинальной модели он действует и так далее.

Для вашего удобства я подготовил репозиторий на HuggingFace содержащий указанный LoRA слой и всё необходимое для его правильной работы, тестовый пример применения можно посмотреть тут.

Шаг 3 - Слияние LoRA слоя и базовой модели

Данный шаг является промежуточным, но он необходим для того чтобы в дальнейшем получить GGML версии модели. Для выполнения процедуры слияния слоёв команда rulm подготовила специальный скрипт под названием convert_to_native.py, но к сожалению он не совместим с ruGPT-3.5, так как оптимизирован под работу с архитектурой LLaMA. В общем пришлось его немного модифицировать. Надо положить его в корень проекта с соответствующим названием, далее создадим файл 3_merge.py со следующим содержанием:

from pathlib import Path
from convert_to_native import convert_to_native

content_dir = Path('.').resolve()
output_dir = content_dir / 'output'
merged_path = output_dir / 'pytorch_model.bin'

convert_to_native(
    model_name=str(output_dir),
    output_path=str(merged_path),
    device='cpu',
    enable_offloading=True
)

assert merged_path.exists()

Исходный код тут.

Скрипт вливает слой LoRA в базовую модель ruGPT-3.5 (загруженную в режиме float32). Для работы скрипта понадобится примерно 60Гб RAM, так как слияние происходит в системной оперативной памяти, что видно по опции device='cpu'.

Давайте запустим его:

python3 3_merge.py

В результате в директории output появится файл pytorch_model.bin, и будет весить примерно 56Гб, по времени процедура слияния занимает примерно 10-15 минут.

Самый любопытный момент в том, что указанный файл уже можно применять для выполнения задач инференса, просто указажите в AutoModelForCausalLM (пакета transformers) путь до папки output.

Пример использования тут.

Шаг 4 - Создание GGML моделей

У нас всё готово для того чтобы начать преобразование pytorch_model.bin в формат GGML, для этого мы будем использовать библиотеку llm-rs-python, которая является python-обёрткой для библиотеки llm, написанной на языке Rust.

Создадим файл 4_ggml.py и наполним его следующим кодом:

from llm_rs.convert import AutoConverter
from llm_rs import AutoQuantizer, QuantizationType, ContainerType
from pathlib import Path

content_dir = Path('.').resolve()
input_dir = content_dir / 'ruGPT-3.5-13B-lora'
output_dir = content_dir / 'output_ggml'

# Convert the model to fp16 format
converted_model = AutoConverter.convert(input_dir, output_dir)

# Quantize the model to different formats
AutoQuantizer.quantize(converted_model, quantization=QuantizationType.Q4_0, container=ContainerType.GGML)
AutoQuantizer.quantize(converted_model, quantization=QuantizationType.Q4_1, container=ContainerType.GGML)
AutoQuantizer.quantize(converted_model, quantization=QuantizationType.Q5_0, container=ContainerType.GGML)
AutoQuantizer.quantize(converted_model, quantization=QuantizationType.Q5_1, container=ContainerType.GGML)
AutoQuantizer.quantize(converted_model, quantization=QuantizationType.Q8_0, container=ContainerType.GGML)

Запустим скрипт, на конвертацию во все форматы понадобится примерно 30 минут:

python3 4_ggml.py

По коду видно, что сначала происходит преобразование модели в формат совместимый с GGML, помимо этого происходит конвертация весов из float32 во float16, затем конвертированная модель сохраняется в директории output_ggml с названием ruGPT-3.5-13B-lora-f16.bin

После преобразования запускается процедура квантизации, по итогу у нас получится 5 версий модели в формате GGML, которые можно запускать например бинарным файлом gpt-2 собранным в рамках проекта ggml или же с помощь llm, или же llm-rs-python и так далее. Все модели будут сохранены в директории output_ggml.

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

Кстати, я подготовил репозиторий на HuggingFace, так что можете их уже пощупать.

Благодарности

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

• Команде Sber AI за оригинальную модель ruGPT-3.5 13B.

• IlyaGusev и команде проекта rulm за датасеты и скрипты используемые для обучения моделей семейства Saiga.

• ggerganov и команде проекта ggml за документацию и исходные коды, которые помогли мне разобраться с тем как правильно запускать модели на процессоре.

• iashchak за репозиторий ruGPT-3.5-13B-ggml на HuggingFace, изучение данного репозитория помогло мне найти проект llm-rs-python.

Заключение

Ну чтож, вот мы и на финишной прямой, надеюсь, эта статья будет полезной для всех, кто интересуется обучением глубоких нейронных сетей и планирует применять модель ruGPT-3.5 13B в своих исследованиях и проектах.

Желаю вам успехов в ваших начинаниях в области машинного обучения!

PS. Решение описанное в данной публикации совместимо с моделями mGPT-13B, так как принципиально их архитектура ничем не отличается от ruGPT-3.5.

Ссылки

• Скрипты для тренировки и все исходники

• ruGPT-3.5 LoRA на HuggingFace

• ruGPT-3.5 GGML на HuggingFace

• Мой Telegram-канал

• Мой блог на Дзене