70M Russian LM (GrokAdamW) + SFT + RAG
Маленькая (71M параметров) decoder-only языковая модель русского языка, обученная с нуля на одной Tesla V100 за ~6.7 часов нашим оптимайзером GrokAdamW. Поверх — инструктивное дообучение (SFT) и RAG (модель остаётся 70M, факты — во внешней базе знаний). Три чекпоинта: lm_ckpt.pt (база), sft_ckpt.pt (чат), rag_ckpt.pt (ответы из контекста).
- Параметры: 71.0M
- Токенов обучения: ~1.05B (≈1.12 эпохи по корпусу 944M)
- Финальная val-перплексия: 25.9
- Корпус: FineWeb-2 (rus_Cyrl) + Русская Википедия (вес 2:1)
- Токенайзер: свой BPE, 32k, byte-level (≈1 токен/слово на русском)
🆕 v4 — диалоговая модель (дистилляция из 7B-учителя) + расширенный претрейн
Цель v4 (на 2×A100): дожать 70M, чтобы она реально вела мультитурн-диалог и при этом сохранила факты из базы. Модель остаётся 71M.
Чекпоинт: chat_rag_ckpt.pt — единая модель, которая И болтает в диалоге, И достаёт факты из RAG.
Что сделано
- Расширенный претрейн (continued pretrain) базы на 3.29B токенов свежего русского (FineWeb-2 + Wiki), bf16, A100. val-ppl на одинаковом срезе: 25.9 → ~23.7 (меньше петель, беглее речь).
- Дистилляция из учителя (SeqKD): локальный Qwen2.5-7B-Instruct сгенерировал 16 473 ответа на разнообразные русские инструкции. 70M-студент учится имитировать стиль 7B-модели (без копирования весов, без внешних API).
- Мультитурн-диалог: 20 734 реальных диалога Saiga в чат-шаблоне
Пользователь:/Ассистент:, лосс только по репликам ассистента — модель учится держать контекст беседы.
- Якорь фактов: 40 000 примеров SberQuAD (
Контекст/Вопрос/Ответ) в том же миксе — чтобы навык извлечения из контекста (RAG) не деградировал.
Итого SFT-микс: 77 207 примеров (19.7M токенов), оптимайзер GrokAdamW, старт с rag_ckpt.pt (а не с нуля — это сохранило факты). По эпохам gap train−val рос 0.0 → 0.12, wd_mult 1.00 → 1.25 (GrokAdamW снова в своей стихии).
Результаты (честно)
| модель | мультитурн-диалог | факты (RAG, 20 вопросов) | val_acc (SFT) |
|---|
| база (претрейн) | нет (петли) | — | — |
sft_ckpt (старый SFT) | только одиночный вопрос | галлюцинации без базы | 0.506 |
rag_ckpt (спец-RAG) | нет | 10/20 | — |
chat_rag_ckpt (v4) | да, держит контекст | 9/20 | 0.593 |
Вывод: v4 добавила настоящий разговорный навык (мультитурн, формат, тон), сохранив почти все факты (9/20 против 10/20 у спец-RAG). Трейд-офф честный: чистая дистилляция «с нуля» давала диалог, но роняла факты до 7/20 — поэтому финальная модель доучена поверх RAG-чекпоинта, что сохранило извлечение из контекста.
Запуск диалога: python3 chat_gen.py --ckpt chat_rag_ckpt.pt --tok ru_tok.json --interactive
Факты из базы: python3 rag_verify_gen.py "Какое самое глубокое озеро в мире?" (использует chat_rag_ckpt.pt + KB).
Пример мультитурн-диалога (temp 0.7): приветствие → «посоветуй книгу» → «что-нибудь полегче» (держит тему) → «поздравь с днём рождения» (связное поздравление). Факты в свободном диалоге модель путает — для фактов используется RAG-режим с базой.
Архитектура
Современный decoder-трансформер (стек уровня Llama/Qwen) с доработками для маленьких LM:
| компонент | значение |
|---|
| dim | 512 |
| слоёв | 17 |
| голов | 8 |
| контекст | 1024 |
| словарь | 32000 |
| норма | RMSNorm (pre-norm) |
| позиции | RoPE |
| FFN | SwiGLU (без bias) |
| внимание | QK-norm + causal SDPA |
| эмбеддинги | tied (вход = выход) |
| софтмакс | z-loss 1e-4 |
| точность | fp16 |
Обучение
- Оптимайзер: GrokAdamW (γ=2.0, weight decay 0.1) — AdamW + decoupled WD + адаптивный WD-импульс по gap = train_acc − val_acc.
- LR: cosine decay, пик 5e-4, warmup 400, пол 10%.
- Батч: 64 эфф. (micro 16 × accum 4), seq 1024 → 64k токенов/шаг.
- grad-clip 1.0; 16000 шагов; ~39.8k токенов/с на V100.
Честно про GrokAdamW на языке
На языковом претрейне плато зубрёжки (train↑/val↓) нет, gap train−val держался ≈ 0.004–0.013 всё обучение, поэтому WD-импульс почти не включался (wd_mult ≈ 1.0). На обильных данных GrokAdamW корректно сводится к AdamW. Его выигрыш проявляется на задачах с резким плато обобщения (напр. модулярная арифметика: 1.63× ускорение гроккинга), а не на языковом претрейне.
Instruction-tuning (SFT)
Поверх базовой модели сделан SFT (инструктивное дообучение), чтобы модель отвечала на вопросы/инструкции, а не просто продолжала текст.
- Датасет:
IlyaGusev/saiga_scored — топовый русский инструктивный набор (на нём учат Saiga/Vikhr). Отфильтровано до сверхкачественного среза: только русский, однотуровые, opus_score ≥ 8 (оценка качества Claude Opus), без regex-мусора → 10 782 пары «вопрос→ответ» (2.9M токенов).
- Формат:
Вопрос: {q}\nОтвет: {a}<eos>. Лосс считается только по ответу (промпт и паддинг маскируются -100). Right-padding безопасен: causal-внимание не «смотрит» на паддинг в конце.
- Оптимайзер: снова GrokAdamW (γ=2.0, wd=0.1), LR 2e-4 cosine, 3 эпохи, лучший чекпоинт по val.
GrokAdamW наконец «в своей стихии»
В отличие от претрейна, на SFT появляется реальный gap train−val (мало примеров, модель переобучается), и адаптивный WD-импульс реально включается:
| эпоха | val_acc | train_acc | gap | wd_mult |
|---|
| 0 | 0.497 | 0.503 | 0.006 | 1.01 |
| 1 (лучшая) | 0.506 | 0.570 | 0.064 | 1.13 |
| 2 | 0.503 | 0.637 | 0.133 | 1.27 |
То есть механизм honest-честно отрабатывает: чем сильнее модель переобучается, тем сильнее автоматически давит weight decay (wd_mult 1.0 → 1.27). Деплоится лучший чекпоинт (эпоха 1).
До / после SFT (одинаковые вопросы)
| вопрос | база (претрейн) | после SFT |
|---|
| Кто написал «Война и мир»? | Тоже. Прислали. Тоже. Прислали… (петля) | Связный ответ в формате ответа (факт неточен — потолок 70M) |
| Назови столицу Франции. | Россия — это временный, временный… (петля) | Париж |
| Поздравление с днём рождения | — | Дорогой [Имя], Поздравляю тебя с днём рождения!… С любовью и лучшими пожеланиями, [Ваше имя] |
База сваливалась в петли на любом вопросе; после SFT модель отвечает в формате ассистента (списки, markdown, краткие ответы где надо, шаблоны писем). Грамотность русского — высокая. Фактическая точность остаётся слабой — это потолок 70M, не источник фактов.
RAG: 70M-модель + своя база знаний (факты не в весах, а во внешней базе)
Главная слабость маленькой LM — факты (их физически некуда положить в 70M). Решение без раздувания модели: retrieval-augmented generation (RAG). Модель остаётся 70M, а знания живут в отдельной базе — и модель учится отвечать из найденного контекста.
- База знаний: 400 000 лид-параграфов русской Википедии (
kb.tar.gz) + индекс BM25 (поиск без нейросети — модель честно остаётся 70M). Заголовок усилен (title×3) + ре-ранкинг по совпадению заголовка с вопросом.
- Дообучение (RAG-SFT):
SberQuAD (45 328 троек контекст+вопрос+ответ из вики) + saiga_scored (10 782, общий чат) = 56 110 примеров. Формат Контекст: …\nВопрос: …\nОтвет:, лосс только по ответу. Оптимайзер — снова GrokAdamW.
- Инференс: вопрос → BM25 достаёт нужный кусок вики → модель отвечает по нему (
rag_gen.py).
Результаты (факты — то, что база/SFT не могли)
| вопрос | retrieved статья | ответ RAG-модели |
|---|
| Кто написал «Война и мир»? | Война и мир | Лев Николаевич Толстой ✓ |
| Столица Франции? | Франция | Париж ✓ |
| Кто написал «Преступление и наказание»? | Преступление и наказание | Достоевский ✓ |
| Когда началась Вторая мировая? | Вторая мировая война | 2 сентября 1939 года ✓ |
| Кто такой Пушкин? | Пушкин, Александр Сергеевич | Русский писатель ✓ |
Для сравнения базовая модель на этих вопросах зацикливалась, а SFT-без-контекста уверенно ошибалась (Война и мир → «Булгаков»). С RAG 70M-модель достаёт правильный факт из контекста.
Честно про ошибки: оставшиеся осечки (ДНК, закон тяготения, «самое глубокое озеро») — это промахи поиска (редиректы вроде ДНК→«Дезоксирибонуклеиновая кислота», суперлативы, лид-параграф без нужного факта), а не модели: когда BM25 достаёт верную статью, ответ верный. Доказательство «заземления» (grounding): на ошибочно найденной статье «Прокляты и убиты» модель честно ответила «Виктор Астафьев» — автор именно той книги.
GrokAdamW на RAG-SFT (импульс снова включается)
| эпоха | val_acc | gap train−val | wd_mult |
|---|
| 0 | 0.509 | 0.021 | 1.04 |
| 1 | 0.529 | 0.081 | 1.16 |
| 2 (деплой) | 0.530 | 0.187 | 1.37 |
RAG v2: база ×5 + гибридный поиск + «внутреннее состояние»
Вторая итерация бьёт в главное узкое место RAG v1 — поиск (модель читает контекст верно, но BM25 иногда достаёт не ту статью / факт не в лид-абзаце).
1. База знаний v2 (kb2.tar.gz): не 400k лид-абзацев, а 2 094 125 пассажей — каждая статья режется на под-куски (до 3 по ~600 символов на границе предложений), так что факты из тела статьи («Байкал — самое глубокое», «восемь планет») тоже попадают в индекс. ~800k статей вики.
2. Гибридный поиск (rag2_gen.py): BM25 (лексика) + плотные эмбеддинги intfloat/multilingual-e5-small (семантика, 2.09M×384, fp16) со слиянием RRF (reciprocal-rank fusion) + бонус за совпадение заголовка. Ловит и точные сущности, и перефразировки. Эмбеддер — только для поиска, генератор честно остаётся 70M.
3. «Внутреннее состояние» (на запрос «чтобы держала смысл») — lm_state.py + rag_state.pt. Механизм (без магии, честно):
- Смысловой якорь: e5-вектор
вопрос+контекст (384-d) проецируется в пространство модели и добавляется к каждому эмбеддингу токена (broadcast-add). Это постоянный «топик-якорь» на каждом шаге. Проекция инициализируется нулём → стартует точь-в-точь как RAG v1 (никакого вреда), учится использовать якорь, только если он помогает.
- Лосс связности: из скрытого состояния на участке ответа читается предсказание смысла (линейный
state_out → 384-d) и притягивается косинусом к e5-вектору ответа. Заставляет остаточный поток нести смысл целиком, а не только следующий токен.
- Доп-параметры (anchor_in + state_out ≈ 0.4M) — только обвязка; ядро-генератор = 71.0M.
Честный итог по «внутреннему состоянию»
Скрытое состояние реально научилось кодировать смысл ответа: cos(state, e5(ответ)) = 0.92 на валидации (против ~0 у необученного). Аддитивный дизайн убрал зацикливание (повторы 0.000). НО прироста фактической точности нет — даже на 1 пункт ниже plain RAG. То есть «держать смысл» получилось технически, но это не повышает интеллект 70M-модели.
Бенчмарк: 20 фактических вопросов (greedy, авто-подсчёт по ключевым словам)
| модель | факты | повторы (bigram-rep, ↓ лучше) |
|---|
| SFT (без базы) | 9/20 | 0.008 |
| RAG v1 (база v2 + гибрид) | 9/20 | 0.000 |
| RAG + внутреннее состояние | 8/20 | 0.000 |
Метрика строгая и шумная (SFT набирает «попадания» галлюцинируя частотные слова). Реальная ценность RAG — заземление и отсутствие петель: Толстой, Париж, «1 сентября 1939», Ньютон, «Евгений Онегин»→Пушкин, Солнце→«одно из звёзд нашей Галактики». Оставшиеся ошибки — промахи поиска (Гагарин→не та статья, «чёрная дыра»→сериал, гора→не та), а не генерации.
Рекомендация: для фактов используйте rag_ckpt.pt с базой v2 и гибридным поиском (rag2_gen.py). rag_state.pt — задокументированный эксперимент «внутреннего состояния»: держит смысл (cos 0.92) и не зацикливается, но по точности не превосходит plain RAG.
RAG v3: «логика на инференсе» — самопроверка модели (это сработало)
«Внутреннее состояние» в весах потолок 70M не двигает. Поэтому логику вынес в процедуру вокруг модели на инференсе — без роста модели и без внешних API. Три приёма проверены численно на том же бенчмарке (eval_verify.py):
- LM-самопроверка пассажей — модель сама оценивает свою уверенность (mean softmax-prob) в ответе по каждому из топ-K найденных пассажей; так 70M-модель судит, какой контекст релевантен.
- Самосогласованность (вотинг) — ответы по разным пассажам голосуют; согласованные усиливают друг друга.
- Заземляющий logit-bias — подкрутка вероятности токенов, реально присутствующих в контексте (экстрактивное «притяжение» к источнику).
Что честно сработало, а что нет
| метод | факты | заметка |
|---|
| RAG v1 (RRF top-3 concat) | 9/20 | baseline |
| VERIFY (per-passage + grounding-bias + vote) | 9/20 | сильный bias скатывался в копирование заголовка («Гагарин.») |
| VERIFY2 (контекст = LM-реранкинг по уверенности) | 8/20 | чинит recall (озеро→Байкал!), но иногда промотит уверенно-неверный пассаж |
| VERIFY3 (объединение RRF-top3 ∪ LM-confidence) | 10/20 | аддитивно: сохраняет попадания baseline + добавляет recall от самопроверки |
Вывод честно: «вбить логику» в 70M-веса не вышло (это потолок ёмкости), но вынос рассуждения на инференс дал реальный измеримый прирост: 10/20 против 9/20, без петель (rep 0.000). Победил аккуратный приём — объединение лексико-семантического поиска с самопроверкой модели (verify3 в eval_verify.py, rag_verify_gen.py для интерактива). Агрессивные приёмы (жёсткий grounding-bias, замена контекста реранкингом) перетасовывают, какие вопросы проходят, но в сумме не помогают. Это и есть честный способ добавить логику там, где параметров не хватает: рассуждение — в процедуре вокруг модели.
Файлы
lm_ckpt.pt — базовый чекпоинт (претрейн): {"model": state_dict, "cfg": {...}, "step": 16000}.sft_ckpt.pt — инструктивный (SFT) чекпоинт (чат без контекста): {"model", "cfg", "epoch", "val_acc"}.rag_ckpt.pt — RAG-чекпоинт (отвечает из контекста, рекомендуется для фактических вопросов).rag_state.pt — RAG + «внутреннее состояние» (эксперимент: якорь + лосс связности, см. RAG v2).kb.tar.gz — база знаний v1: 400k пассажей вики (passages.jsonl) + BM25-индекс.kb2.tar.gz — база знаний v2: 2.09M пассажей (passages.jsonl) + BM25-индекс + плотные эмбеддинги e5 (embs.f16.npy).rag2_gen.py — RAG v2 инференс (гибридный поиск BM25+e5+RRF).rag_verify_gen.py — RAG v3 инференс (самопроверка verify3: RRF ∪ LM-confidence; лучший по фактам, 10/20).rag_state_gen.py — инференс с «внутренним состоянием» (якорь e5).lm_state.py — модель с якорем + головой связности (LMState).kb_build2.py, embed_kb.py, train_state.py, eval_factual.py, eval_verify.py — пайплайн v2/v3 (сбор базы, эмбеддинги, обучение состояния, бенчмарки).ru_tok.json — токенайзер (HF tokenizers).lm_model.py — определение модели (LM, LMConfig).gen.py — генерация (продолжение текста, база).sft_gen.py — генерация ответов (SFT, repetition penalty).rag_gen.py — RAG-инференс (BM25-поиск + ответ из контекста).kb_build.py, kb_reindex.py, sft_rag_data.py, sft_train.py, sft_data.py — пайплайн воспроизведения.
Запуск
1import torch, torch.nn.functional as F
2from tokenizers import Tokenizer
3from lm_model import LM, LMConfig
4
5tok = Tokenizer.from_file("ru_tok.json")
6ck = torch.load("lm_ckpt.pt", map_location="cpu")
7model = LM(LMConfig(**ck["cfg"]))
8model.load_state_dict(ck["model"]); model.eval().cuda()
9
10@torch.no_grad()
11def gen(prompt, n=110, temp=0.8, top_k=40):
12 ids = [tok.token_to_id("<bos>")] + tok.encode(prompt).ids
13 x = torch.tensor([ids], device="cuda")
14 for _ in range(n):
15 logits = model(x[:, -1024:])[:, -1, :].float() / temp
16 v, _ = torch.topk(logits, top_k)
17 logits[logits < v[:, [-1]]] = -float("inf")
18 nxt = torch.multinomial(F.softmax(logits, -1), 1)
19 x = torch.cat([x, nxt], 1)
20 if nxt.item() == tok.token_to_id("<eos>"): break
21 return tok.decode(x[0].tolist())
22
23print(gen("Москва — столица"))
Чат-режим (SFT)
1# та же модель, грузим sft_ckpt.pt; промпт в формате "Вопрос: ...\nОтвет:"
2ck = torch.load("sft_ckpt.pt", map_location="cpu")
3model = LM(LMConfig(**ck["cfg"])); model.load_state_dict(ck["model"]); model.eval().cuda()
4
5@torch.no_grad()
6def ask(q, n=180, temp=0.6, top_k=40, rep=1.2):
7 ids = [tok.token_to_id("<bos>")] + tok.encode(f"Вопрос: {q}\nОтвет:").ids
8 x = torch.tensor([ids], device="cuda"); start = x.shape[1]
9 for _ in range(n):
10 logits = model(x[:, -1024:])[:, -1, :].float()
11 for t in set(x[0].tolist()): logits[0, t] /= rep
12 logits /= temp
13 v, _ = torch.topk(logits, top_k); logits[logits < v[:, [-1]]] = -float("inf")
14 nxt = torch.multinomial(F.softmax(logits, -1), 1); x = torch.cat([x, nxt], 1)
15 if nxt.item() == tok.token_to_id("<eos>"): break
16 return tok.decode(x[0, start:].tolist()).strip()
17
18print(ask("Назови столицу Франции.")) # -> Париж
RAG-режим (70M + база знаний)
1tar xzf kb.tar.gz # -> kb/passages.jsonl + kb/bm25_index/
2pip install bm25s PyStemmer
3python3 rag_gen.py rag_ckpt.pt # BM25-поиск по вики -> ответ из контекста
Внутри: вопрос → BM25 достаёт лид-параграф нужной статьи → промпт Контекст: …\nВопрос: …\nОтвет: → модель отвечает по контексту.
RAG v2-режим (база 2M + гибридный поиск)
1tar xzf kb2.tar.gz # -> kb2/passages.jsonl + kb2/bm25_index/ + kb2/embs.f16.npy
2pip install bm25s PyStemmer transformers
3python3 rag2_gen.py rag_ckpt.pt # BM25 + e5-эмбеддинги + RRF -> ответ из контекста
4# с «внутренним состоянием»:
5python3 rag_state_gen.py # грузит rag_state.pt + lm_state.py, добавляет e5-якорь
Внутри: вопрос → BM25 (топ-100) + e5-косинус (топ-100) → RRF-слияние + бонус заголовка → топ-3 пассажа → промпт Контекст: …\nВопрос: …\nОтвет:. Для rag_state.pt к эмбеддингам токенов добавляется e5-якорь вопроса+контекста.
RAG v3-режим (самопроверка — лучший по фактам, 10/20)
1tar xzf kb2.tar.gz
2pip install bm25s PyStemmer transformers tokenizers
3python3 rag_verify_gen.py "Какое самое глубокое озеро в мире?" # -> Байкал
Внутри (verify3): достаём топ-8 пассажей → для каждого 70M-модель сама оценивает уверенность ответа → контекст = объединение RRF-топ-3 и пассажей с наибольшей уверенностью модели → финальный ответ по этому контексту. «Логика» вынесена в процедуру вокруг модели (модель судит свой контекст), веса не меняются.
Примеры (temp 0.8, top-k 40)
Москва — столица, в котором проживает около 3 тыс. человек… масштабная работа по внедрению системы водоснабжения… В настоящее время в столице насчитывается 2 тыс. км системы водоснабжения.
Наука изучает… Биография Родился в 1927 году в Москве в семье рабочих. В 1939 году окончил Московский институт инженеров железнодорожного транспорта…
Ограничения
Грамматика и морфология русского — отличные; тема держится на 2–4 предложения. Глобальная когерентность и фактичность — слабые (фактические галлюцинации, иногда самоповторы). Это потолок модели 70M на ~1B токенов — не источник фактов. Качество ограничено размером модели, бюджетом токенов и железом (V100, fp16).
