Update README.md

2c4de01 verified 13 days ago

18.9 kB

language:
  - ru
tags:
  - wiki
  - ru_wiki
  - wikipedia

SunboxDevLLM - Русскоязычная языковая модель

Модель доступна для использования: chat.sunbox.dev

Репозиторий: GitHub

Decoder-only Transformer языковая модель, построенная полностью с нуля — BPE-токенизатор, архитектура модели, цикл обучения и генерация — на Python + PyTorch (~2500 строк кода).

Архитектура

Компонент	Решение	Обоснование
Нормализация	RMSNorm (pre-norm)	Стандарт LLaMA/Mistral, ~10-15% быстрее LayerNorm
Позиции	RoPE (split + cat, без complex)	Лучшая экстраполяция длины, совместимость с MPS/AMP
FFN	SwiGLU (3 проекции)	Выше качество при том же вычислительном бюджете
Attention	F.scaled_dot_product_attention	Flash Attention на CUDA, math fallback на MPS
Генерация	KV cache (prefill + decode)	O(S) на новый токен вместо O(S²)
Веса	Weight tying (embed ↔ lm_head)	Экономия ~24.6M параметров
AMP	autocast + GradScaler	CUDA и MPS (PyTorch 2.2+)

Конфигурации модели

Вариант	d_model	heads	layers	ff_dim	Параметры
Full	768	12	12	2048	~110M

Установка

# Клонировать репозиторий и перейти в корень проекта (имя папки может отличаться)
cd llm

# Создать виртуальное окружение
python3 -m venv .venv
source .venv/bin/activate

# Установить зависимости
pip install -r requirements.txt

Чекпоинты весов (models/*.pt) не коммитятся (см. .gitignore). Для примеров ниже положите файлы base_step_15000.pt и chat_step_500.pt в каталог models/ или передайте свой путь в --checkpoint. Для BPE нужен data/vocab.json из того же обучения (или --vocab), если путь из чекпоинта на этой машине не существует.

Зависимости

Пакет	Обязательный	Назначение
`torch>=2.2.0`	да	PyTorch (CUDA / MPS / CPU)
`numpy>=1.20`	да	Memmap для pretokenized-корпусов
`regex>=2023.0`	да	Unicode-aware pre-tokenization (GPT-2 паттерн)
`transformers>=4.20`	нет	Адаптер HuggingFace-токенизатора (`--tokenizer`)
`tokenizers>=0.15`	нет	Быстрое BPE-обучение через Rust (`train_bpe_hf.py`)
`gensim>=4.0`	нет	Парсинг дампа Wikipedia (`data/wiki_download.py`)

Быстрый старт

Полный пайплайн одной командой

bash run_train.sh

Скрипт выполнит: обучение BPE-токенизатора → pre-training на Wikipedia → fine-tuning на чат-корпусе.

Пошаговый запуск

1. Обучение BPE-токенизатора

# Собственная реализация (чистый Python, ~20-50x быстрее наивного алгоритма)
python train_bpe.py --corpus data/wiki_text/corpus.txt --vocab-size 32000

# Или через HuggingFace tokenizers (Rust, ~100-1000x быстрее на больших корпусах)
python train_bpe_hf.py --corpus data/wiki_text/corpus.txt --vocab-size 32000

2. Pretokenization корпуса (для больших датасетов)

Для корпусов >1 GB рекомендуется pretokenization — после неё обучение стартует мгновенно:

python pretokenize_corpus.py \
  --corpus data/wiki_text/corpus_15gb.txt \
  --output data/corpus_15gb.bin \
  --workers 8

3. Pre-training

# На текстовом корпусе (токенизация при старте)
python train.py --corpus data/wiki_text/corpus.txt \
  --batch-size 32 --max-steps 100000 --save-every 5000

# На pretokenized-корпусе (мгновенный старт)
python train.py --pretokenized data/corpus_15gb.bin \
  --batch-size 64 --grad-accum 4 --max-steps 30000

# С ранней остановкой по perplexity
python train.py --pretokenized data/corpus_15gb.bin \
  --stop-at-perplexity 10.0

4. Fine-tuning на чат-корпусе

# --init-from загружает только веса модели (свежий optimizer/scheduler)
python train.py --corpus data/wiki_text/corpus_chat.txt --chat \
  --init-from models/base_step_15000.pt \
  --lr 1e-4 --batch-size 16 --max-steps 2000 --save-every 250

# --resume загружает полное состояние (продолжение прерванного обучения)
python train.py --corpus data/wiki_text/corpus_chat.txt --chat \
  --resume models/chat_step_500.pt \
  --max-steps 3000

Генерация текста

Режим	temperature	repetition_penalty	top_p	Комментарий
Текст (pre-trained)	0.3	1.2	—	Максимально связные дополнения, минимум шума
Чат (fine-tuned)	0.5	1.1	—	Лучший баланс точности и разнообразия
Универсальный (default)	0.5	1.15	—	Разумный компромисс для обоих режимов
Креативный	0.7	1.15	0.9	Больше разнообразия, выше риск галлюцинаций

Одиночный prompt

# Streaming (токены выводятся по одному)
python generate.py --checkpoint models/chat_step_500.pt \
  --prompt "Москва — столица" --stream

# Полный ответ сразу (без --stream)
python generate.py --checkpoint models/chat_step_500.pt \
  --prompt "Москва — столица" --max-tokens 100

Prompt из файла

python generate.py --checkpoint models/chat_step_500.pt \
  --prompt-file my_prompt.txt --stream

Интерактивный чат

# Базовый режим (каждый вопрос независимый)
python generate.py --checkpoint models/chat_step_500.pt \
  --chat --temperature 0.5 --repetition-penalty 1.1

# С сохранением контекста (sliding window по границам реплик)
python generate.py --checkpoint models/chat_step_500.pt \
  --chat --keep-history

Пример сессии:

=== Chat Mode (type 'quit' to exit) ===

You: Что такое машинное обучение?
Assistant: Машинное обучение — раздел ИИ: алгоритмы, которые улучшают своё поведение
на основе данных без явного программирования каждого шага.

You: Что такое фотосинтез?
Assistant: Фотосинтез — процесс образования органических веществ из CO₂ и воды
на свету при участии хлорофилла. Выделяется кислород.

You: quit

Параметры генерации

Параметр	По умолчанию	Описание
`--temperature`	0.5	Температура сэмплирования (ниже = детерминированнее)
`--top-k`	—	Ограничение по top-k токенам
`--top-p`	—	Nucleus sampling (top-p)
`--max-tokens`	256	Максимум новых токенов
`--repetition-penalty`	1.15	Штраф за повторения (1.0 = выкл, рекомендуется 1.1-1.2)
`--stream`	off	Потоковый вывод токенов
`--chat`	off	Интерактивный чат-режим
`--keep-history`	off	Сохранение контекста диалога между репликами
`--vocab`	auto	Путь к vocab.json (если чекпоинт обучен на другой машине)

Параметры обучения

Параметр CLI	Конфиг	По умолчанию	Описание
`--max-steps`	`max_steps`	23000	Максимум шагов
`--batch-size`	`batch_size`	32	Размер батча
`--lr`	`learning_rate`	6e-4	Learning rate
`--grad-accum`	`grad_accum_steps`	4	Gradient accumulation
`--warmup`	`warmup_steps`	2000	Шаги warmup
`--save-every`	`save_every_steps`	2000	Частота сохранения
`--seq-len`	`max_seq_len`	1024	Длина контекста
`--no-amp`	`use_amp`	True	Отключить AMP
`--stop-at-loss`	`stop_at_loss`	—	Ранняя остановка по loss
`--stop-at-perplexity`	`stop_at_perplexity`	—	Ранняя остановка по val perplexity

LR scheduler: linear warmup → cosine annealing (single cycle).

Подготовка данных

Скачивание Wikipedia

python data/wiki_download.py --max-mb 500 --corpus-file data/wiki_text/corpus.txt

Очистка корпуса

Удаляет блоки EasyTimeline, галереи, лишние пустые строки:

python clean_corpus.py --input data/wiki_text/corpus.txt \
  --output data/wiki_text/corpus_cleaned.txt --remove-gallery

# Только статистика без записи
python clean_corpus.py --dry-run

Обрезка по размеру

python trim_corpus.py --input data/wiki_text/corpus_cleaned.txt \
  --output data/wiki_text/corpus_500mb.txt --mb 500

Генерация чат-корпуса

python data/wiki_text/generate_chat_corpus.py

Создаёт corpus_chat.txt из пула Q&A-пар на основе seed-данных (столицы, события, персоны и т.д.).

Структура проекта

llm/
├── config.py                   # ModelConfig, TrainConfig, GenerateConfig, PathsConfig
├── train.py                    # CLI обучения
├── generate.py                 # CLI генерации и интерактивного чата
├── train_bpe.py                # Обучение BPE (собственная реализация)
├── train_bpe_hf.py             # Обучение BPE (HuggingFace Rust backend)
├── pretokenize_corpus.py       # Pretokenization корпуса в .bin (multiprocessing)
├── clean_corpus.py             # Очистка wiki-корпуса
├── trim_corpus.py              # Обрезка корпуса по размеру
├── run_train.sh                # Полный пайплайн одной командой
│
├── model/
│   ├── transformer.py          # TransformerLM (блоки + lm_head + weight tying)
│   ├── block.py                # DecoderBlock (pre-norm) + SwiGLUFFN
│   ├── attention.py            # MultiHeadAttention (RoPE + SDPA + KV cache)
│   ├── rope.py                 # Rotary positional embeddings
│   ├── norm.py                 # RMSNorm
│   └── kv_cache.py             # KV cache для генерации
│
├── tokenizer/
│   ├── vocab.py                # Словарь (tuple pair keys, без overflow)
│   ├── bpe.py                  # BPETokenizer (pre-tokenize + merge ranks + LRU)
│   ├── bpe_train_fast.py       # Быстрое обучение BPE (doubly-linked list + heap)
│   ├── pretokenizer.py         # GPT-2 regex word splitter
│   └── hf_adapter.py           # Адаптер HuggingFace-токенизатора
│
├── data/
│   ├── wiki_dataset.py         # WikiChunkDataset + PretokenizedDataset (memmap)
│   ├── chat_dataset.py         # ChatChunkDataset (маскирование loss на user-части)
│   ├── wiki_download.py        # Скачивание/парсинг дампа Wikipedia
│   └── wiki_text/              # Корпуса и генераторы чат-данных
│
├── training/
│   ├── trainer.py              # AMP, gradient accumulation, валидация, логирование
│   ├── checkpoint.py           # Сохранение/загрузка полного состояния
│   └── scheduler.py            # Warmup + CosineAnnealingWarmRestarts
│
├── models/                     # *.pt в .gitignore — веса кладутся локально
│   ├── base_step_15000.pt      # Pre-trained checkpoint (Wikipedia 15GB)
│   └── chat_step_500.pt        # Fine-tuned checkpoint (chat corpus)
│
├── inference/
│   └── generator.py            # KV cache prefill/decode, streaming, sampling
│
└── utils/
    ├── device.py               # Детекция CUDA/MPS/CPU и платформенные настройки
    └── log_setup.py            # Конфигурация логирования

Совместимость

Платформа	Поддержка	Особенности
CUDA GPU (A100, V100, T4)	полная	Flash Attention, GradScaler, pin_memory, num_workers=4
Apple MPS (M1/M2/M3/M4)	полная	SDPA math fallback, unified memory, num_workers=0
CPU	базовая	bfloat16 autocast где поддерживается

Отличия от оригинального llm

CUDA-совместимость: единый get_device() с приоритетом CUDA > MPS > CPU
BPE без overflow: tuple-ключи (id_a, id_b) вместо bit-shift (vocab > 65536)
Быстрый BPE encode: pre-tokenization (GPT-2 regex) + merge rankings + LRU cache
RoPE вместо sinusoidal: лучшая экстраполяция, без отдельного embedding-слоя
RMSNorm вместо LayerNorm: быстрее, соответствует LLaMA/Mistral
SwiGLU вместо GELU FFN: выше качество при том же compute
KV cache при генерации: O(S) на токен вместо O(S²) полного re-forward
Weight tying: embed и lm_head разделяют веса, экономия ~24.6M параметров
Полное состояние чекпоинтов: модель + optimizer + scheduler + scaler + step
Grad norm до clipping: точный мониторинг здоровья градиентов
Pre-tokenized datasets: O(1) __getitem__, корректное маскирование loss для чата
GPT-2 инициализация: N(0, 0.02) + residual scaling 0.02 / sqrt(2N)

SunboxDev
/

SunBoxDev-Chat-Ru-110M