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
Генерация текста
Рекомендованные параметры
Параметры подобраны тестированием на base_step_15000 (pre-trained) и chat_step_500 (chat fine-tuned)
на 8 различных промптах по 9 конфигурациям каждый. Значения по умолчанию в коде (GenerateConfig)
установлены как универсальный компромисс: temperature=0.5, repetition_penalty=1.15.
| Режим | 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 | Больше разнообразия, выше риск галлюцинаций |
Примеры запуска с оптимальными параметрами:
# Дополнение текста (pre-trained checkpoint)
python generate.py --checkpoint models/base_step_15000.pt \
--prompt "Москва — столица" --temperature 0.3 --repetition-penalty 1.2 --stream
# Чат (fine-tuned checkpoint) — оптимальные параметры
python generate.py --checkpoint models/chat_step_500.pt \
--chat --temperature 0.5 --repetition-penalty 1.1
# Без явных параметров — используются defaults из GenerateConfig
python generate.py --checkpoint models/chat_step_500.pt \
--prompt "Искусственный интеллект — это" --stream
Одиночный 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 scaling0.02 / sqrt(2N)