Как я обучил модель, которая понимает русский лучше GPT 3.5 Turbo

В этой статье я расскажу, как я смог обучить модель, которая превзошла GPT 3.5 Turbo на русскоязычной части MT-Bench. Также я рассмотрю новую конфигурацию для обучения на двух графических процессорах параллельно с помощью accelerate и deepspeed.

Особенный интерес представляет мой датасет для обучения. Он получен из сабсета мультиязычных промтов набора lightblue/tagengo-gpt4 на русском, английском и китайском, всего 10 тысяч примеров, сгенерированных с помощью GPT-4o. Это в 8 раз меньше, чем исходный набор Tagengo, но обученная на последнем Suzume, как показали бенчмарки, лишь очень незначительно превосходит мою модель на ru_mt_bench, а на англоязычном бенче и вовсе уступает ей. Это значит, что я в разы сэкономил на GPU за счет более высокого качества данных, полученных с помощью GPT-4o. 

Я использовал скрипт для получения ответов по заданным промптам. Для генерации русскоязычной выборки я изменил часть скрипта, чтобы выбрать все промпты на русском из Tagengo (8K примеров), так как основной фокус при обучении модели был на русском языке. 

В итоге я получил датасет ruslandev/tagengo-rus-gpt-4o и приступил к обучению.

Для этого я создал виртуальную машину с NVIDIA H100, используя сервис immers.cloud. Для достижения наилучших результатов по instruction-following (что проверяется на MT-Bench) я взял в качестве исходной модели meta-llama/Meta-Llama-3-8B-Instruct. Именно на ней обучена модель Suzume, у которой высокая оценка на MT Bench. Предыдущие эксперименты показали, что базовая Llama-3 8B,  а особенно ее четырехбитная версия для QLoRA — unsloth/llama-3-8b-bnb-4bit — значительно отстает по оценкам бенчмарка.

В этот раз я обучил параллельно на двух GPU, для этого я задал новую конфигурацию моей виртуальной машины — две NVIDIA A100.

Я воспользовался инструментом axolotl, который позволяет быстро сконфигурировать и запустить сессию обучения.

Мой конфиг axolotl здесь.

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

accelerate launch -m axolotl.cli.train config.yaml

Accelerate — это Huggingface библиотека для распределенного обучения.

axolotl запустил два параллельных процесса с шардами модели для каждого из двух GPU. Обучение на одну эпоху длилось около часа, итоговый train loss — 0.8.

Результат превзошел мои ожидания — третье место в mt_bench:

Моя модель превзошла llama-3-8b-instruct и большинство версий Suzume, кроме самой сильной из них. Это на англоязычном бенчмарке.

Теперь — результат ru_mt_bench:

Моя модель получила оценку 8.12, немного не дотянув до Suzume и превосходя gpt-3.5-turbo, у которой 7.94.

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

Я не увеличивал англоязычную выборку в своем датасете, там всего тысяча примеров на английском, а англоязычный MT Bench неожиданно показал среднюю оценку 8 баллов. Это значит, что добавление большего количества качественных мультиязычных данных повышает общее качество модели, а не только ее показатели для этого конкретного языка. Эта эффект уже был показан в статье Петера Девина — Tagengo: A Multilingual Chat Dataset

Я очень рад, что мне удалось наблюдать реализацию этой идеи на практике. Мой датасет, веса модели, а также файлы в формате GGUF опубликованы в моем Huggingface аккаунте.

Модель

GGUF

Датасет