Дообучение ruGPT-3.5 13B с LoRA08.10.2023 20:46

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

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

Сгенерировано Kandinsky 2.2

Сгенерировано Kandinsky 2.2

В данной публикации мы разберем каждый этап обучения модели, начиная от подготовки данных и заканчивая конвертацией в формат 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 от Гигабайт, но даже её памяти не хватало для запуска модели. Поэтому я стал искать различные способы её хоть как-то уместить в память карточки, по ходу дела узнал про библиотеку bitsanbytes, которая расширяет функционал библиотеки 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.

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

К слову сказать, оригинальная 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)

По коду видно, что сначала происходит преобразование модели в формат совместимый с 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 в своих исследованиях и проектах.

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

Ссылки

© Habrahabr.ru