README.md

---
language:
- ru
tags:
- wiki
- ru_wiki
- wikipedia
---
# SunboxDevLLM - Русскоязычная языковая модель

Модель доступна для использования: **[chat.sunbox.dev](https://chat.sunbox.dev)**

Репозиторий: **[GitHub](https://github.com/it-efrem/llm)**

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     |

## Установка

```bash
# Клонировать репозиторий и перейти в корень проекта (имя папки может отличаться)
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
bash run_train.sh
```

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

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

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

```bash
# Собственная реализация (чистый 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 — после неё обучение стартует мгновенно:

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

#### 3. Pre-training

```bash
# На текстовом корпусе (токенизация при старте)
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 на чат-корпусе

```bash
# --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  | Больше разнообразия, выше риск галлюцинаций  |

Примеры запуска с оптимальными параметрами:

```bash
# Дополнение текста (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

```bash
# 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 из файла

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

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

```bash
# Базовый режим (каждый вопрос независимый)
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

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

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

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

```bash
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
```

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

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

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

```bash
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

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