Как я собрал русскоязычную GPT на 112 млн параметров с нуля на PyTorch: претрейн, SFT, Telegram-бот

Когда смотришь на современную LLM, легко решить, что всё сводится к вызову API. Но многие разработчики, которые часто используют LLM в работе или в своих проектах, рано или поздно хотят действительно разобраться, как они работают. К тому же - кто не смотрел «Железного человека»? По сути самый лучший способ реально разобраться это сделать что-то свое. Чаще всего возможности позволяют сделать что-то маленькое, зато свое и работающее.

Так появился jarvis - компактный decoder-only transformer для русского текста и диалогов. В проекте нет готового AutoModelForCausalLM: я сам реализовал архитектуру, обучил токенизатор, собрал пайплайн подготовки данных, добавил претрейн и SFT, KV-cache для генерации, консольный чат, Telegram-бота и простой бенчмарк. Готовый SFT-checkpoint и соответствующий токенизатор опубликованы на Hugging Face.

Что получилось в итоге

Вот состав проекта.

Структура репозитория довольно прямолинейна:

model/                 # RoPE, RMSNorm, attention, блоки, Transformer
prepare_data.py        # получение и очистка данных
train_tokenizer.py     # обучение BPE-токенизатора
tokenize_corpus.py     # бинарные token shards и маски для SFT
dataset.py             # streaming dataset через numpy.memmap
train.py               # претрейн/SFT, DDP
chat.py, generate.py   # локальный инференс
bot.py                 # Telegram-бот
tools/bench_bpc.py     # оценка BPC
selftest.py            # синтетический сквозной тест

Дальше разберу несколько своих решений.

Модель

Основной пресет выглядит так:

Внутри нет nn.Transformer. Блок собран из двух pre-norm residual-веток:

class TransformerBlock(nn.Module):
    def forward(self, x, cache=None):
        attn_out, new_cache = self.attn(self.ln1(x), cache)
        x = x + attn_out
        x = x + self.ff(self.ln2(x))
        return x, new_cache

Attention и RoPE

Для Q, K и V используется одна линейная проекция без bias, затем тензор разбивается на три части и на головы. RoPE применяется к Q и K, а собственно вычисление внимания делегировано оптимизированному torch.nn.functional.scaled_dot_product_attention:

qkv = self.qkv_proj(x)
qkv = qkv.reshape(batch_size, seq_len, 3, self.n_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]

q, k = self.rope.rotate(q, k, offset)
attn_output = F.scaled_dot_product_attention(
    q, k, v,
    is_causal=(seq_len == k.shape[2]),
    dropout_p=self.dropout if self.training else 0.0,
)

Почему не написать attention целиком руками через матрицу QKᵀ, маску и softmax? На первых шагах так и хочется сделать: код банально проще и нагляднее. Но в рабочем проекте выгоднее сохранить прозрачность архитектуры и одновременно использовать хорошо оптимизированное ядро PyTorch. У меня всё ещё свои QKV-проекции, свои RoPE, своя работа с cache и causal-логикой; не приходится платить за медленную Python-реализацию матричных операций.

RMSNorm, SwiGLU и tied embeddings

Вместо LayerNorm здесь RMSNorm - нормализация по среднему квадрату без центрирования. MLP - SwiGLU: один линейный слой даёт gate, второй значения, дальше они перемножаются после SiLU, а третий слой возвращает размерность к d_model.

Ещё одно полезное решение - связать входные эмбеддинги и выходную проекцию в словарь:

self.token_embedding = nn.Embedding(cfg.vocab_size, cfg.d_model)
self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
self.lm_head.weight = self.token_embedding.weight

Это сокращает число независимых параметров и является обычным приёмом для компактной языковой модели.

Токенизатор

В проекте используется свой byte-level BPE с размером словаря 16 384. Byte-level подход не ломается на незнакомых символах: в крайнем случае текст представляется байтами. При этом BPE всё равно учит частые фрагменты русских слов и пунктуации как более крупные токены.

Токенизатор обучается на смеси обычных текстов и диалогов. В train_tokenizer.py на текстовые корпуса приходится 60% символьного бюджета, на чаты 40%. Это важно: модель должна видеть не только литературный русский, но и тот формат, в котором потом будет получать запросы.

Помимо [UNK], в словарь добавляются четыре служебных токена:

<|endoftext|>
<|user|>
<|assistant|>
<|pad|>

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

<|user|>Привет!<|assistant|>Привет! Чем помочь?<|endoftext|>

Есть один нюанс. В текущем формате у system нет собственного специального токена: системный текст присоединяется к первому сообщению пользователя. Это упрощает формат и нормализацию разнородных датасетов.

Почему SFT нельзя считать по всему диалогу

Для претрейна естественная цель проста: предсказывать следующий токен везде. С диалоговым дообучением есть тонкость. Если считать loss по всей последовательности, модель одновременно учится продолжать и пользователя, и ассистента. В интерфейсе же нас интересует второе.

Поэтому для диалогов рядом с токенами хранится маска. Токены пользователя получают вес 0, токены ответа ассистента 1. Упрощённо это выглядит так:

<|user|>   Как дела?   <|assistant|>   Всё хорошо.   <|endoftext|>
    0           0             0               1              1

В коде маска строится по служебным токенам: после <|assistant|> включается, при <|user|> выключается, завершающий <|endoftext|> тоже относится к ответу. Во время SFT loss умножается на эту маску:

losses = F.cross_entropy(
    logits.view(-1, logits.shape[-1]),
    y.view(-1),
    reduction="none",
)
weights = mask.view(-1).float()
loss = (losses * weights).sum() / weights.sum().clamp_min(1.0)

Эта деталь превращает набор диалогов в датасет для обучения отвечающей модели.

Данные: три типа текстов

Я разделил данные на три источника.

Скрипт не складывает всё подряд в один мешок. Для веба и gazeta есть фильтры по длине текста и доле кириллицы. Для диалогов нормализуются названия ролей, склеиваются подряд идущие реплики одной роли, отбрасываются невалидные примеры. Повторы отсеиваются.

По умолчанию prepare_data.py собирает до миллиона веб-документов и до 150 тысяч текстов gazeta; для russian_instructions_2 есть отдельный лимит в 120 тысяч примеров. Сырые корпуса не коммитятся в репозиторий - их можно получить заново из источников:

python prepare_data.py --what all
python train_tokenizer.py
python tokenize_corpus.py --source all

Бинарные шарды и memmap

После токенизации получаются бинарные шарды по 1 млн uint16-токенов. Для чатов дополнительно записывается файл масок. На обучении их не нужно целиком загружать в RAM: TokenStreamDataset открывает shard через numpy.memmap, выбирает источник согласно весам, затем случайный шард и случайное окно нужной длины.

Такой дизайн не претендует на что-то сложное, но у него есть практические плюсы: простой формат, быстрый случайный доступ, небольшой расход памяти и отсутствие зависимости от тяжёлой базы данных.

Две фазы обучения

Конфигурация намеренно разделена на претрейнинг и SFT.

На первой фазе модель учит статистику языка и широкого текста. Диалоговые примеры уже присутствуют в смеси, но loss считается на всех токенах. На второй фазе доля диалогов резко возрастает, а для них остаются только токены ответов ассистента.

Запуск выглядит так:

python train.py --preset main --phase pretrain

python train.py --preset main --phase sft \
  --resume checkpoints/pretrain/latest.pt

Основной пресет использует context length 768, micro-batch 24 и gradient accumulation 2. На одной GPU это 36 864 токена на один шаг; при DDP объём масштабируется с числом GPU.

Что делает training loop, кроме loss.backward()

В train.py есть несколько вещей, которые легко упустить в первом эксперименте, а потом очень не хватает при длинном запуске:

• AdamW с betas=(0.9, 0.95) и weight decay;

• линейный warmup, затем cosine decay learning rate;

• gradient clipping до нормы 1.0;

• автоматический выбор fp16/bf16 AMP для CUDA;

• DDP, если доступно несколько GPU;

• ограничение длительности сессии по умолчанию 11 часами;

• периодические checkpoint каждые 15 минут;

• --resume auto для продолжения последней сессии;

• отдельный inference.pt с весами fp16;

• валидация и сохранение примеров на фиксированных prompts.

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

Локальная генерация и KV-cache

Авторегрессионная генерация на каждом новом токене снова прогоняет через модель весь промпт. На коротком примере это терпимо, но с ростом истории становится всё заметнее.

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

Поверх этого в generate.py есть базовые:

• temperature;

• top_k;

• top_p (nucleus sampling);

• repetition penalty;

• остановка по <|endoftext|> и следующему <|user|>.

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

git clone https://github.com/lubenkov/jarvis.git
cd jarvis

python -m venv .venv
# Активируйте виртуальное окружение
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
python -m pip install --upgrade huggingface_hub

Сначала нужно установить подходящую для своей системы сборку PyTorch - особенно это важно для CUDA. Затем скачиваем checkpoint и tokenizer вместе:

python - <<'PY'
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="Eragon1/jarvis-112m",
    local_dir="checkpoints/hf-jarvis-112m",
    allow_patterns=["jarvis-112m-sft-fp16.pt", "tokenizer.json"],
)
PY

И запускаем чат:

python chat.py \
  --ckpt checkpoints/hf-jarvis-112m/jarvis-112m-sft-fp16.pt \
  --tokenizer checkpoints/hf-jarvis-112m/tokenizer.json

В чате доступны exit, /reset и /temp 0.7. Для одного запроса есть generate.py.

Важное ограничение: checkpoint это обычный PyTorch .pt в формате проекта, а не готовый артефакт transformers. Команда вида AutoModelForCausalLM.from_pretrained() здесь не сработает; для загрузки нужен generate.load_model() из репозитория и совпадающий токенизатор.

Зачем здесь Telegram-бот

Консоль хороша для разработки, но не даёт почувствовать модель в обычном диалоге. Поэтому я добавил необязательный bot.py.

python -m pip install -r requirements-bot.txt
cp .env.example .env
# Укажите TELEGRAM_BOT_TOKEN в .env

python bot.py \
  --ckpt checkpoints/hf-jarvis-112m/jarvis-112m-sft-fp16.pt \
  --tokenizer checkpoints/hf-jarvis-112m/tokenizer.json

Бот хранит историю отдельно для каждого чата, обрезает её под контекст, умеет /reset, /temp и /stats. Но важнее не набор команд, а защита от банальных проблем одиночного локального инференса: есть TTL для неактивных чатов, lock на конкретный чат, лимит активных диалогов, ограниченная очередь обновлений, минимальный интервал между запросами и семафор на число одновременных генераций. По умолчанию генерация одна за раз - это реалистичнее, чем делать вид, будто одна GPU обслужит бесконечный поток пользователей.

Что показывает BPC-бенчмарк

Для базовой автоматической оценки я использовал bits per character (BPC) на отложенных фрагментах по 600 000 символов. Меньше = лучше.

Результат не самый громкий, зато он полезен. После SFT модель немного лучше предсказывает диалоговый текст (0.92 → 0.91 BPC) но хуже справляется с вебом и статьями. То есть специализация действительно смещает модель к диалогам, а не улучшает всё сразу.

Сравнение с ruGPT3small стоит читать осторожно. У моделей разные токенизаторы и разные окна контекста, а BPC отвечает только на вопрос насколько хорошо модель предсказывает следующий символ в этом тексте. Эта метрика ничего напрямую не говорит о полезности ответа, фактической точности, безопасности или умении следовать инструкции. Поэтому я не называю таблицу лидербордом и не использую её как доказательство превосходства модели.

Код измерения лежит в tools/bench_bpc.py, а полный JSON с числами в results/bpc.json.

Проверки

В репозитории есть selftest.py. Он намеренно не скачивает реальные корпуса и опубликованные веса, вместо этого создаёт синтетический маленький корпус и проходит полный короткий маршрут.

python selftest.py

Скрипт проверяет:

1. прямой проход модели и совпадение логитов при по-токенной генерации с KV-cache;

2. обучение небольшого токенизатора;

3. токенизацию корпуса в train/validation;

4. маленький претрейн-запуск;

5. продолжение обучения;

6. переход от чекпоинта претрейна к SFT;

7. генерацию из сохранённого inference.pt.

Такой тест не гарантирует качество модели, но хорошо ловит поломки на стыках компонентов. Для проекта с несколькими CLI-скриптами это намного полезнее, чем проверять каждый файл только вручную.

Ограничения

У небольших моделей есть своя ниша, но у них нет суперспособностей. Для jarvis я бы особенно выделил следующее:

• 112 млн параметров — это очень мало по меркам современных универсальных LLM;

• окно контекста 768 токенов быстро заканчивается;

• модель ориентирована прежде всего на русский язык;

• ответы могут быть неверными, повторяющимися, предвзятыми или небезопасными;

• модель не прошла полноценную оценку безопасности и не должна быть единственным источником решения в медицинских, юридических, финансовых и других важных задачах;

• лицензия кода и чекпоинта Apache-2.0, но лицензии и условия исходных датасетов нужно проверять отдельно;

Именно поэтому карточка модели на Hugging Face не обещает невозможного: там опубликованы архитектура, способ запуска, источники данных, BPC-результаты и ограничения, но нет выдуманных сведений о железе, точном фактически пройденном числе токенов или человеческом уровне качества.

Итог

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

Исходники: GitHub
Веса и model card: Hugging Face
Лицензия: Apache-2.0

Буду рад техническим замечаниям по данным и оценкам.