Несколько дней к ряду я занимался реставрацией легаси модели ai-forever/rugpt3xl, это классическая языковая модель от SberDevices на 1.3B параметров, крошка по современным меркам, на которой сберовцы обкатывали свои научные наработки аж в далёком 2021м году. Подробнее о ней можно почитать в статье “A family of pretrained transformer language models for Russian” на Google Scholar.

Да, она foundation, то есть умеет только продолжать текст, не может выполнять инструкции или работать в режиме чата. Но обучена она на корпусе русского языка и этот самый русский генерит очень бодро. У неё есть две примечательные особенности: её обучали с нуля, архитектура представляет собой глубокую модификацию GPT-2.

Превьюшку сгенерировал через ChatGPT
Превьюшку сгенерировал через ChatGPT

Предыстория

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

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

Проблема была в том, что оригинальный чекпоинт суть сырой mp_rank_00_model_states.pt, чтобы его загрузить, нужен полный стек Megatron-LM, DeepSpeed, apex, и всё это до кучи завязано на древние версии PyTorch 1.7 и transformers 3.5, короче говоря, запустить её в 2026 году “как есть” задача нетривиальная.

Но решил попробовать ещё один раз, благо думаю где мне знаний будет мало кодовый агент поможет, поднял старые заметки, склонировал всё что может понадобиться для работы и написал подробные спецификации, предварительную работу делал Cursor, а я уже тестировал и пинал фантазии ИИ в нужную сторону.

Конвертация в HuggingFace формат

Задача агента была простая, для начала он должен был изучить исходники ai-forever/ru-gpts, исходники современной библиотеки transformers, веса и конфиги модели deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct в качестве примера того, как на HuggingFace можно закидывать кастомные классы моделей и конфигов, чтобы всё работало корректно с trust_remote_code=True.

Моя основная цель была в том, чтобы получить полностью рабочую модель, идентичную оригинальной на столько на сколько это возможно, которую можно было бы запустить через transformers, и при этом чтобы её можно было обучать, хоть через LoRA, хоть через SFT-тренеры из поставки transformers.

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

Вот как выглядит оригинальная структура модели из чекпоинта Megatron-LM:

word_embeddings.weight                 [50264, 2048]
position_embeddings.weight             [2048, 2048]
transformer.layers.{0..23}
  +-- input_layernorm.weight           [2048]
  +-- input_layernorm.bias             [2048]
  +-- attention
  |     +-- query_key_value.weight     [6144, 2048]  <- сшитые QKV
  |     +-- query_key_value.bias       [6144]
  |     +-- dense.weight               [2048, 2048]
  |     +-- dense.bias                 [2048]
  +-- post_attention_layernorm.weight  [2048]
  +-- post_attention_layernorm.bias    [2048]
  +-- mlp
        +-- dense_h_to_4h.weight       [8192, 2048]
        +-- dense_h_to_4h.bias         [8192]
        +-- dense_4h_to_h.weight       [2048, 8192]
        +-- dense_4h_to_h.bias         [2048]
transformer.final_layernorm.weight     [2048]
transformer.final_layernorm.bias       [2048]

А вот так выглядит структура GPT2 моделей:

model.embed_tokens.weight              [50264, 2048]
model.embed_positions.weight           [2048, 2048]
model.layers.{0..23}
  +-- input_layernorm.weight           [2048]
  +-- input_layernorm.bias             [2048]
  +-- self_attn
  |     +-- q_proj.weight              [2048, 2048]  <- отдельно Q
  |     +-- q_proj.bias                [2048]
  |     +-- k_proj.weight              [2048, 2048]  <- отдельно K
  |     +-- k_proj.bias                [2048]
  |     +-- v_proj.weight              [2048, 2048]  <- отдельно V
  |     +-- v_proj.bias                [2048]
  |     +-- o_proj.weight              [2048, 2048]
  |     +-- o_proj.bias                [2048]
  +-- post_attention_layernorm.weight  [2048]
  +-- post_attention_layernorm.bias    [2048]
  +-- mlp
        +-- up_proj.weight             [8192, 2048]
        +-- up_proj.bias               [8192]
        +-- down_proj.weight           [2048, 8192]
        +-- down_proj.bias             [2048]
model.norm.weight                      [2048]
model.norm.bias                        [2048]
lm_head.weight                         [50264, 2048]

Главное отличие в том, что сшитая QKV проекция размерностью [6144, 2048] разбивается на три отдельных линейных слоя Q, K, V размерностью [2048, 2048] каждый, плюс явный lm_head, который нужно будет скопировать из весов эмбеддингов.

Вот полная таблица маппинга:

Megatron-LM (оригинал)

HuggingFace (конвертированная)

word_embeddings.weight

model.embed_tokens.weight

position_embeddings.weight

model.embed_positions.weight

transformer.layers.{i}.input_layernorm.*

model.layers.{i}.input_layernorm.*

transformer.layers.{i}.attention.query_key_value.weight

model.layers.{i}.self_attn.{q,k,v}_proj.weight

transformer.layers.{i}.attention.query_key_value.bias

model.layers.{i}.self_attn.{q,k,v}_proj.bias

transformer.layers.{i}.attention.dense.*

model.layers.{i}.self_attn.o_proj.*

transformer.layers.{i}.post_attention_layernorm.*

model.layers.{i}.post_attention_layernorm.*

transformer.layers.{i}.mlp.dense_h_to_4h.*

model.layers.{i}.mlp.up_proj.*

transformer.layers.{i}.mlp.dense_4h_to_h.*

model.layers.{i}.mlp.down_proj.*

transformer.final_layernorm.*

model.norm.*

-

lm_head.weight (копия model.embed_tokens.weight)

Зная маппинг выполнить конвертацию не составлаяет уже большого труда, получился скрипт convert.py, который берёт оригинальный чекпоинт и превращает его в HuggingFace-модель с safetensors весами.

Но конвертация весов это только полдела, далее потребовалось написать кастомные классы модели и конфигурации, которые transformers смог бы загрузить.

В оригинальном репозитории ru-gpts есть класс RuGPT3XL это толстая обёртка вокруг мегатроновского GPT3Model, завёрнутого ещё и в FP16_Module. Чтобы просто загрузить модель, from_pretrained первым делом поднимает torch.distributed процесс-группу, инициализирует mpu (model parallelism utils), скачивает веса и deepspeed-конфиг, и только потом собирает модель через setup_model. KV-кеша нет, без CUDA и torch.distributed ничего не работает, красивый артефакт своей эпохи, но для современного пользователя всё это просто мёртвый груз.

Новые классы писались с нуля по примерам transformers, иерархия такая:

  • RuGPT3XLConfig - наследник PretrainedConfig из трансформерс со всеми параметрами модели (vocab_size, hidden_size, num_layers и так далее)

  • RuGPT3XLAttention - multi-head attention с отдельными Q/K/V проекциями и поддержкой DynamicCache из трансформерс

  • RuGPT3XMLP - блок персептрона (up_proj -> GELU -> down_proj)

  • RuGPT3XLDecoderLayer - слой декодера (LayerNorm -> Attention -> LayerNorm -> MLP)

  • RuGPT3XLModel - базовый трансформер модели (эмбеддинги + 24 слоя + финальный LayerNorm)

  • RuGPT3XLForCausalLM - обёртка конструктор с lm_head

Принципиальные отличия от оригинала:

  • Никаких зависимостей от Megatron-LM, DeepSpeed, apex, mpu, torch.distributed и прочей древности

  • Типовой forward() с сигнатурой, которую ожидает любой тренер, хоть LoRA/Pert, хоть SFTTrainer

  • Результаты возвращаются в CausalLMOutputWithPast из трансформерс вместо кастомного ModelOutput

  • Полноценный KV-кеш через трансформеровский DynamicCache для быстрой генерации

  • Поддержка gradient checkpointing на этапе обучения для оптимального управления памятью во время обучения

  • Ну и самое главное, моделька работает на CPU, GPU, через device_map="auto" без обязательного CUDA

По сути математика внутри осталась идентичной - тот же Pre-LayerNorm, та же GELU, те же размерности, тот же токенизатор, но вот обвязка вокруг модельки стала современной.

Веса залил на HuggingFace: evilfreelancer/ruGPT3XL.

Тестирование

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

python generate.py --model_path ./ruGPT3XL --dtype float16

На выходе три дефолтных промпта, и моделька вполне связный текст пишет - “Москва - столица”, “Искусственный интеллект - это”, “В далёком космосе”. Видно, что русский язык она знает хорошо, фразы строит грамотно, контекст держит, для foundation-модельки на 1.3B параметров прямиком из 2021 года это более чем достойно.

Дальше захотелось потыкать модельку посерьёзнее. Написал manual_test.py - скрипт на 24 промпта разного характера, от свободного продолжения текста (“Однажды в студёную зимнюю пору”, “Для приготовления борща нужно”) до простых вопрос-ответов через chat template (“Вопрос: Какая столица России? Ответ:”), гонял на RTX 4090 48Гб во float16, тут результаты.

Пара интересных наблюдений по результатам:

  • Средняя скорость генерации - 66.7 t/s на RTX 4090 (float16, batch_size=1)

  • На коротких промптах моделька часто останавливается рано (15-30 токенов), на длинных уверенно генерит до max_new_tokens

  • Вопрос-ответ через chat template работает, но ответы иногда… творческие. На “Кто написал Войну и мир?” модель уверенно ответила “Пушкин” (ну, foundation-модель, чтож поделать)

  • Хорошо справляется с продолжением фактического текста - рецепт борща, физика, история

  • На “Сколько планет в Солнечной системе?” ответила “Пять”, а потом ушла в рассуждения про цивилизации Млечного Пути

Ну и напоследок прогнал модельку через MERA, это такой открытый бенчмарк для оценки русскоязычных моделей, 23 задачи разного типа (логика, математика, знания о мире, код, этика, рассуждения). Загрузил результаты на mera.a-ai.ru, общий скор получился 0.198 (пока ещё на модерации).

Результаты MERA
Результаты MERA

Для понимания контекста - это base-модель 2021 года на 1.3B параметров, без instruction tuning, без RLHF, без всего того, что сейчас считается обязательным. Она умеет только продолжать текст. И тем не менее на некоторых задачах показывает вменяемые результаты:

  • PARus (здравый смысл) - 0.500

  • ruHateSpeech - 0.558

  • BPS (код/математика) - 0.528

  • RWSD (рассуждения) - 0.488

  • ruTiE (диалоговый контекст) - 0.502

  • ruMMLU - 0.252 (при random baseline ~0.25 для 4 вариантов ответа, модель на грани случайного угадывания, что для base-модели ожидаемо)

А вот математика (SimpleAr - 0.012, ruModAr - 0.001) и генерация кода (ruHumanEval, ruCodeEval - 0/0/0) ожидаемо в ноль, ведь модель и не учили решать подобные задачи.

Конвертация в GGUF

Но я не остановился на достигнутом. Захотелось развить успех и конвертировать модельку в GGUF формат для llama.cpp. Благодаря тому, что я потратил время на конвертацию в формат transformers, реализовать поддержку в llama.cpp удалось очень быстро. Там всего-навсего нужно было прокачать Python-конвертер convert_hf_to_gguf.py - модель по структуре похожа на GPT-2, поэтому патч объединяет отдельные q_proj, k_proj, v_proj обратно в единый QKV тензор, как того ожидает LLM_ARCH_GPT2.

Отправил PR #21011 в llama.cpp с наработками (кстати второй мой PR принятый в этот проект), прошёл ревью, были мелкие правки по хешу токенизатора, в итоге PR смерджили, а я смог конвертировать и квантовать модельки.

GGUF залил на HuggingFace: evilfreelancer/ruGPT3XL-GGUF.

Запуск инференса прямо в терминале:

llama-cli -hf evilfreelancer/ruGPT3XL-GGUF:Q4_K_M

Или поднять локальный OpenAI-совместимый сервер с веб-интерфейсом:

llama-server -hf evilfreelancer/ruGPT3XL-GGUF:Q4_K_M

Ну и под конец залил модельку на Ollama, там теперь собрано всё семейство ruGPT-3 в одном месте, все четыре размера - small (125M), medium (356M), large (760M) и xl (1.3B):

ollama run evilfreelancer/rugpt3:small
ollama run evilfreelancer/rugpt3:medium
ollama run evilfreelancer/rugpt3:large
ollama run evilfreelancer/rugpt3:xl

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

Ссылки

Послесловие

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

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