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 — пайплайн воспроизведения.

Запуск

import torch, torch.nn.functional as F
from tokenizers import Tokenizer
from lm_model import LM, LMConfig

tok = Tokenizer.from_file("ru_tok.json")
ck = torch.load("lm_ckpt.pt", map_location="cpu")
model = LM(LMConfig(**ck["cfg"]))
model.load_state_dict(ck["model"]); model.eval().cuda()

@torch.no_grad()
def gen(prompt, n=110, temp=0.8, top_k=40):
    ids = [tok.token_to_id("<bos>")] + tok.encode(prompt).ids
    x = torch.tensor([ids], device="cuda")
    for _ in range(n):
        logits = model(x[:, -1024:])[:, -1, :].float() / temp
        v, _ = torch.topk(logits, top_k)
        logits[logits < v[:, [-1]]] = -float("inf")
        nxt = torch.multinomial(F.softmax(logits, -1), 1)
        x = torch.cat([x, nxt], 1)
        if nxt.item() == tok.token_to_id("<eos>"): break
    return tok.decode(x[0].tolist())

print(gen("Москва — столица"))

Чат-режим (SFT)

# та же модель, грузим sft_ckpt.pt; промпт в формате "Вопрос: ...\nОтвет:"
ck = torch.load("sft_ckpt.pt", map_location="cpu")
model = LM(LMConfig(**ck["cfg"])); model.load_state_dict(ck["model"]); model.eval().cuda()

@torch.no_grad()
def ask(q, n=180, temp=0.6, top_k=40, rep=1.2):
    ids = [tok.token_to_id("<bos>")] + tok.encode(f"Вопрос: {q}\nОтвет:").ids
    x = torch.tensor([ids], device="cuda"); start = x.shape[1]
    for _ in range(n):
        logits = model(x[:, -1024:])[:, -1, :].float()
        for t in set(x[0].tolist()): logits[0, t] /= rep
        logits /= temp
        v, _ = torch.topk(logits, top_k); logits[logits < v[:, [-1]]] = -float("inf")
        nxt = torch.multinomial(F.softmax(logits, -1), 1); x = torch.cat([x, nxt], 1)
        if nxt.item() == tok.token_to_id("<eos>"): break
    return tok.decode(x[0, start:].tolist()).strip()

print(ask("Назови столицу Франции."))   # -> Париж

RAG-режим (70M + база знаний)

tar xzf kb.tar.gz                 # -> kb/passages.jsonl + kb/bm25_index/
pip install bm25s PyStemmer
python3 rag_gen.py rag_ckpt.pt    # BM25-поиск по вики -> ответ из контекста

Внутри: вопрос → BM25 достаёт лид-параграф нужной статьи → промпт Контекст: …\nВопрос: …\nОтвет: → модель отвечает по контексту.

RAG v2-режим (база 2M + гибридный поиск)

tar xzf kb2.tar.gz               # -> kb2/passages.jsonl + kb2/bm25_index/ + kb2/embs.f16.npy
pip install bm25s PyStemmer transformers
python3 rag2_gen.py rag_ckpt.pt  # BM25 + e5-эмбеддинги + RRF -> ответ из контекста
# с «внутренним состоянием»:
python3 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)

tar xzf kb2.tar.gz
pip install bm25s PyStemmer transformers tokenizers
python3 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).

Downloads last month: -; Downloads are not tracked for this model. How to track