Практикум — претрейн маленькой модели

Пошаговый проект — русские данные, свой токенизатор, модель на 4–26M параметров с RMSNorm, RoPE, SwiGLU и GQA, обучение, графики loss, базовые линии и генерация; ориентиры уровней Simple/Medium/Strong/Boss.

УРОКstage-5/5-6-minimind/02-pretrain-project.mdx4 чУрок 2 из 2

После урока вы сможете

Подготовить русский корпус — очистить, дедуплицировать, обучить токенизатор и упаковать токены в бинарный поток
Собрать и обучить модель на 4–26M параметров с современным блоком на CPU или одной GPU
Читать кривые loss, сравнивать модель с n-граммными базовыми линиями и считать BPB
Планировать эксперименты и бюджет вычислений по законам масштабирования

Зачем это нужно

Вы прочитали, как устроены законы масштабирования, данные и MiniMind. Пора собрать всё в одном проекте: своя маленькая LLM, которая пишет по-русски. Не дообучение чужой модели, а предобучение с нуля — от сырого текста до генерации.

Такой проект учит вещам, которые не видны в статьях: как выглядит здоровая кривая loss, сколько на самом деле занимает обучение, почему модель повторяет одну фразу и что делать, когда loss стал nan. Модель будет маленькой (4–26M параметров), но конвейер — тот же, что у DeepSeek и Qwen, только без тысяч видеокарт.

Ниже — план из шести шагов и рабочий код, который мы проверили запуском. Код написан в духе nanoGPT и MiniMind: всё на чистом PyTorch, около 300 строк на весь проект.

План проекта

Бюджет. Выберите размер модели и число токенов калькулятором из урока 5.1.1.
Данные. Соберите русский корпус, очистите и дедуплицируйте его (код из урока 5.1.2).
Токенизатор. Обучите byte-level BPE на своих данных и упакуйте токены в бинарный файл.
Модель. Соберите трансформер с RMSNorm, RoPE, SwiGLU и GQA.
Обучение. Запустите цикл с AdamW, разогревом и косинусом, логируйте train и val loss.
Оценка. Сравните с n-граммными базовыми линиями, постройте графики, посмотрите на генерацию.

Структура папки:

tiny-llm/
├── corpus.jsonl     # {"text": "..."} — по документу на строку
├── prepare.py       # токенизатор + train.bin / val.bin
├── model.py         # модель
├── train.py         # обучение → ckpt.pt, log.json
├── baselines.py     # униграммы и биграммы
├── plot.py          # loss.png
└── sample.py        # генерация

Зависимости: torch, numpy, matplotlib и tokenizers (pip install tokenizers).

Шаг 0. Бюджет

Сначала посчитайте, во что вы ввязываетесь. Три конфигурации модели, которые мы будем использовать:

Конфигурация	$d_{model}$ × слоёв	Головы Q / KV	FFN	Словарь	Параметров
XS — отладка	256 × 4	4 / 2	704	4096 (6400)	4.0M (4.6M)
S — CPU, вечер	384 × 6	6 / 2	1024	6400	11.9M
M — одна GPU	512 × 8	8 / 2	1408	6400	25.8M

По Chinchilla каждой нужно около 20 токенов на параметр. Сколько это стоит, зависит от того, сколько FLOPs в секунду реально выдаёт ваше железо. Для оценки возьмём условные 0.2 TFLOPS у процессора и 20 TFLOPS у видеокарты:

Конфигурация	Токенов по Chinchilla	$C = 6ND$	CPU, 0.2 TFLOPS	GPU, 20 TFLOPS
XS	80M	$1.9 \cdot 10^{15}$	≈3 ч	≈2 мин
S	238M	$1.7 \cdot 10^{16}$	≈24 ч	≈15 мин
M	517M	$8.0 \cdot 10^{16}$	≈4.6 сут	≈1.1 ч

Вывод: на CPU реалистично обучить S на десятках миллионов токенов, то есть меньше оптимума. Для обучения это нормально: модель всё равно заговорит, просто хуже, чем могла бы. На видеокарте можно позволить себе M «по Chinchilla» и даже больше.

Шаг 1. Данные

Где взять русский текст. Три надёжных варианта:

FineWeb-2, подмножество rus_Cyrl — очищенный русский веб, лицензия ODC-By. Скачивается потоком через библиотеку datasets: load_dataset("HuggingFaceFW/fineweb-2", name="rus_Cyrl", split="train", streaming=True) — возьмите первые 100–500 тыс. документов.
Русская Википедия — дампы на dumps.wikimedia.org/ruwiki/, лицензия CC BY-SA. Плотный энциклопедический текст.
Классика в общественном достоянии — Пушкин, Толстой, Чехов. Хороший литературный язык, но мало объёма и устаревшая лексика. Проверяйте, что берёте сами тексты, а не охраняемые издания с комментариями.

Совет из TinyStories. Маленькой модели проще выучить простой язык. Модели размером около 10M параметров, обученные на простых сказках, пишут связные истории, а на «взрослом» вебе — обрывки. Если хотите, чтобы модель S писала связно, соберите корпус попроще: детские рассказы, простые статьи, или сгенерируйте простые сказки большой моделью (проверьте условия использования её API).

Очистка. Прогоните документы через quality_check и deduplicate из практики урока 5.1.2. Для корпуса на сотни тысяч документов попарная дедупликация будет медленной — используйте LSH (полосы подписи как ключи словаря) или библиотеку datasketch. Выброшенное смотрите глазами: на Википедии правила вроде «доли строк со знаком конца предложения» могут резать таблицы и списки, и пороги придётся подвинуть.

Формат. corpus.jsonl, по документу на строку: {"text": "..."}, как в MiniMind. Разбиение на train и val — по документам, а не по токенам: иначе куски одного текста попадут в обе части, и val loss будет нечестным.

Шаг 2. Токенизатор и упаковка

Токенизатор MiniMind для русского не подходит (урок 5.6.1), поэтому учим свой. Код повторяет подход MiniMind: byte-level BPE из библиотеки tokenizers. Можно взять и свою реализацию BPE из ДЗ s4-hw01-bpe-tokenizer.

import json, random
import numpy as np
from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers

VOCAB = 6400
SPECIAL = ["<|endoftext|>", "<|im_start|>", "<|im_end|>"]

docs = [json.loads(line)["text"] for line in open("corpus.jsonl", encoding="utf-8")]
random.seed(0)
random.shuffle(docs)
n_val = max(1, len(docs) // 100)                  # 1% документов — в валидацию
val_docs, train_docs = docs[:n_val], docs[n_val:]

# 1) Byte-level BPE, как в MiniMind (trainer/train_tokenizer.py)
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tok.decoder = decoders.ByteLevel()
trainer = trainers.BpeTrainer(vocab_size=VOCAB, special_tokens=SPECIAL,
                              initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator(train_docs, trainer=trainer)
tok.save("tokenizer.json")

# 2) Упаковка: все документы — в один поток токенов через <|endoftext|>
eot = tok.token_to_id("<|endoftext|>")
def encode(texts):
    ids = []
    for enc in tok.encode_batch(texts):
        ids.extend(enc.ids)
        ids.append(eot)
    return np.array(ids, dtype=np.uint16)       # словарь меньше 65 536 → хватает uint16

meta = {"vocab_size": tok.get_vocab_size()}
for name, part in [("train", train_docs), ("val", val_docs)]:
    arr = encode(part)
    arr.tofile(f"{name}.bin")
    meta[f"{name}_tokens"] = len(arr)
    meta[f"{name}_bytes"] = sum(len(t.encode("utf-8")) for t in part)
    print(f"{name}: {len(part):,} документов, {len(arr):,} токенов")

meta["bytes_per_token"] = meta["val_bytes"] / meta["val_tokens"]
json.dump(meta, open("meta.json", "w"))
print(f"словарь: {meta['vocab_size']}, байт на токен: {meta['bytes_per_token']:.2f}")
ids = tok.encode("Мудрая сова нашла странную книгу.").ids
print([tok.decode([i]) for i in ids])

На маленьком тестовом корпусе последняя строка печатает что-то вроде:

['Мудрая', ' сова', ' на', 'ш', 'л', 'а', ' странную', ' книгу', '.']

Частые слова стали целыми токенами, редкие («нашла») разобраны на куски. Две вещи, на которые стоит посмотреть:

Выводите токены через tok.decode([i]), а не .tokens. Byte-level BPE хранит токены как байты, и кириллица в .tokens выглядит как кракозябры (кракозябры): каждая буква — это два байта. Если токен содержит только половину буквы, decode покажет на его месте � — это нормально, в составе текста байты соберутся обратно.
Размер словаря. Для русского корпуса в десятки миллионов токенов разумно 4096–8192. Словарь в 32 000 дал бы модели S $32\,000 \cdot 384 \approx 12$ M параметров только на эмбеддинги — больше, чем все её слои вместе (≈9.4M).

Упаковка вместо паддинга. MiniMind дополняет каждый документ паддингом до max_seq_len. Мы делаем, как nanoGPT: склеиваем все документы в один поток через <|endoftext|> и при обучении вырезаем случайные окна. Это упаковка последовательностей: ни одной пустой позиции, все FLOPs идут в дело. Цена — окно может начаться с середины документа, но модель быстро учится, что после <|endoftext|> начинается новый текст.

Шаг 3. Модель

Это блок из урока 4.5.3 в коде: RMSNorm вместо LayerNorm, RoPE вместо обучаемых позиций, SwiGLU вместо GELU-MLP, GQA вместо полного MHA, никаких bias и общие веса эмбеддингов и LM-головы. Внимание считает F.scaled_dot_product_attention — на GPU он сам выберет FlashAttention.

import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class Config:
    vocab_size: int = 6400
    dim: int = 384          # d_model
    n_layers: int = 6
    n_heads: int = 6        # головы запросов
    n_kv_heads: int = 2     # головы ключей/значений (GQA: 3 запроса на 1 KV-голову)
    ffn_dim: int = 1024     # скрытый размер SwiGLU (≈ 8/3 · dim, кратно 64)
    max_seq_len: int = 256
    rope_theta: float = 10_000.0
    dropout: float = 0.0


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        # нормируем по среднему квадрату (без вычитания среднего и без bias)
        rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return (x.float() * rms).type_as(x) * self.weight


def rope_cache(head_dim, max_len, theta):
    inv_freq = 1.0 / theta ** (torch.arange(0, head_dim, 2).float() / head_dim)
    angles = torch.outer(torch.arange(max_len).float(), inv_freq)   # (T, head_dim/2)
    angles = torch.cat([angles, angles], dim=-1)                       # (T, head_dim)
    return angles.cos(), angles.sin()


def apply_rope(x, cos, sin):
    # x: (B, H, T, hd); поворачиваем пары (x_i, x_{i+hd/2})
    x1, x2 = x.chunk(2, dim=-1)
    rotated = torch.cat([-x2, x1], dim=-1)
    return x * cos + rotated * sin


class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_heads, self.n_kv = cfg.n_heads, cfg.n_kv_heads
        self.hd = cfg.dim // cfg.n_heads
        self.q = nn.Linear(cfg.dim, cfg.n_heads * self.hd, bias=False)
        self.k = nn.Linear(cfg.dim, cfg.n_kv_heads * self.hd, bias=False)
        self.v = nn.Linear(cfg.dim, cfg.n_kv_heads * self.hd, bias=False)
        self.o = nn.Linear(cfg.n_heads * self.hd, cfg.dim, bias=False)
        self.dropout = cfg.dropout

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)
        k = self.k(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
        v = self.v(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        rep = self.n_heads // self.n_kv                     # GQA: размножаем KV-головы
        k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                           dropout_p=self.dropout if self.training else 0.0)
        return self.o(y.transpose(1, 2).reshape(B, T, -1))


class SwiGLU(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.gate = nn.Linear(cfg.dim, cfg.ffn_dim, bias=False)
        self.up = nn.Linear(cfg.dim, cfg.ffn_dim, bias=False)
        self.down = nn.Linear(cfg.ffn_dim, cfg.dim, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1, self.attn = RMSNorm(cfg.dim), Attention(cfg)
        self.norm2, self.ffn = RMSNorm(cfg.dim), SwiGLU(cfg)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm1(x), cos, sin)    # Pre-Norm + residual
        return x + self.ffn(self.norm2(x))


class TinyLLM(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
        self.norm = RMSNorm(cfg.dim)
        self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight          # weight tying
        cos, sin = rope_cache(cfg.dim // cfg.n_heads, cfg.max_seq_len, cfg.rope_theta)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)
        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        T = idx.shape[1]
        x = self.tok_emb(idx)
        cos, sin = self.cos[:T], self.sin[:T]
        for blk in self.blocks:
            x = blk(x, cos, sin)
        logits = self.lm_head(self.norm(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=50):
        for _ in range(max_new_tokens):
            logits, _ = self(idx[:, -self.cfg.max_seq_len:])
            logits = logits[:, -1, :] / temperature
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = -float("inf")
            nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
            idx = torch.cat([idx, nxt], dim=1)
        return idx


def count_params(model):
    # общие (tied) веса считаются один раз
    return sum(p.numel() for p in model.parameters())

Проверка здравого смысла до обучения — размеры и начальный loss на случайных данных:

import math, torch
from model import Config, TinyLLM, count_params

torch.manual_seed(0)
configs = {
    "XS (отладка на CPU)": Config(vocab_size=4096, dim=256, n_layers=4, n_heads=4, n_kv_heads=2, ffn_dim=704),
    "S  (CPU, вечер)":     Config(vocab_size=6400, dim=384, n_layers=6, n_heads=6, n_kv_heads=2, ffn_dim=1024),
    "M  (1 GPU)":          Config(vocab_size=6400, dim=512, n_layers=8, n_heads=8, n_kv_heads=2, ffn_dim=1408),
}
for name, cfg in configs.items():
    print(f"{name}: {count_params(TinyLLM(cfg)) / 1e6:.2f}M параметров")

cfg = configs["S  (CPU, вечер)"]
model = TinyLLM(cfg)
x = torch.randint(0, cfg.vocab_size, (4, 128))
y = torch.randint(0, cfg.vocab_size, (4, 128))
logits, loss = model(x, y)
print("logits:", tuple(logits.shape))
print(f"начальный loss = {loss.item():.3f}, ln V = {math.log(cfg.vocab_size):.3f}")

XS (отладка на CPU): 4.00M параметров
S  (CPU, вечер): 11.90M параметров
M  (1 GPU): 25.83M параметров
logits: (4, 128, 6400)
начальный loss = 8.847, ln V = 8.764

Начальный loss близок к $\ln V$ — тот самый приём из урока 1.3.2: свежая модель предсказывает почти равномерное распределение. Конфигурация M по размерам совпадает с minimind2-small (26M): 8 слоёв по 512, 8 голов запросов и 2 KV-головы, словарь 6400 — только токенизатор свой, русский.

Шаг 4. Обучение

Цикл в стиле nanoGPT. Данные читаются через np.memmap: файл на диске открывается как массив NumPy, а числа подгружаются только тогда, когда к ним обращаются. Поэтому файл на гигабайты не нужно грузить в память целиком.

import json, math, time
import numpy as np
import torch
from model import Config, TinyLLM, count_params

meta = json.load(open("meta.json"))
# ── гиперпараметры (конфигурация XS: 4–5M параметров в зависимости от словаря) ──
cfg = Config(vocab_size=meta["vocab_size"], dim=256, n_layers=4, n_heads=4, n_kv_heads=2,
             ffn_dim=704, max_seq_len=256)
batch_size, accum = 16, 2                  # токенов на шаг: 16 · 2 · 256 = 8192
max_steps, warmup = 2000, 100
lr_max, lr_min = 1e-3, 1e-4
eval_every, eval_iters = 100, 20
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(42)

train = np.memmap("train.bin", dtype=np.uint16, mode="r")
val = np.memmap("val.bin", dtype=np.uint16, mode="r")

def get_batch(data):
    ix = torch.randint(len(data) - cfg.max_seq_len, (batch_size,))   # начала 0 … len−T−1: верхняя граница не входит
    x = torch.stack([torch.from_numpy(data[i:i + cfg.max_seq_len].astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy(data[i + 1:i + 1 + cfg.max_seq_len].astype(np.int64)) for i in ix])
    return x.to(device), y.to(device)

def lr_at(step):                            # линейный разогрев + косинус до lr_min
    if step < warmup:
        return lr_max * (step + 1) / warmup
    t = min(1.0, (step - warmup) / max(1, max_steps - warmup))
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * t))

@torch.no_grad()
def evaluate(data):
    model.eval()
    losses = [model(*get_batch(data))[1].item() for _ in range(eval_iters)]
    model.train()
    return sum(losses) / len(losses)

model = TinyLLM(cfg).to(device)
n_params = count_params(model)
print(f"параметров: {n_params / 1e6:.2f}M, устройство: {device}, ln V = {math.log(cfg.vocab_size):.2f}")
decay = [p for p in model.parameters() if p.dim() >= 2]       # матрицы — с weight decay
no_decay = [p for p in model.parameters() if p.dim() < 2]     # веса RMSNorm — без
opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1},
                         {"params": no_decay, "weight_decay": 0.0}],
                        lr=lr_max, betas=(0.9, 0.95))
use_bf16 = device == "cuda" and torch.cuda.is_bf16_supported()
tokens_per_step = batch_size * accum * cfg.max_seq_len

log, best_val, t0 = [], float("inf"), time.time()
for step in range(max_steps + 1):
    if step % eval_every == 0:
        tr, va = evaluate(train), evaluate(val)
        elapsed = time.time() - t0
        tok_s = step * tokens_per_step / max(elapsed, 1e-9)
        bpb = va / math.log(2) / meta["bytes_per_token"]
        log.append({"step": step, "tokens": step * tokens_per_step, "train": tr, "val": va, "bpb": bpb})
        print(f"шаг {step:5d} | train {tr:.3f} | val {va:.3f} | BPB {bpb:.3f} | "
              f"{tok_s:,.0f} ток/с ≈ {6 * n_params * tok_s / 1e12:.3f} TFLOPS")
        if va < best_val:
            best_val = va
            torch.save({"cfg": cfg.__dict__, "model": model.state_dict()}, "ckpt.pt")
        json.dump(log, open("log.json", "w"))
    if step == max_steps:
        break
    for g in opt.param_groups:
        g["lr"] = lr_at(step)
    for _ in range(accum):                                    # накопление градиентов
        x, y = get_batch(train)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_bf16):
            _, loss = model(x, y)
        (loss / accum).backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    opt.zero_grad(set_to_none=True)

Почему именно так:

AdamW с $\beta_2 = 0.95$ (а не 0.999 по умолчанию) — стандарт для LLM: оптимизатор быстрее реагирует на всплески градиента. Weight decay 0.1 — только для матриц; веса норм и одномерные параметры не штрафуют.
Разогрев и косинус. Первые 100 шагов learning rate растёт от нуля: в начале обучения градиенты большие и шумные. Затем косинус до $0.1 \cdot \eta_{max}$ , как в MiniMind.
Обрезка градиента по норме 1.0 страхует от редких огромных шагов, которые ломают обучение.
Накопление градиентов: эффективный батч batch_size · accum при памяти на batch_size.
BF16 включается только на GPU, которые его поддерживают. На CPU обучение идёт в FP32.
Лучший чекпоинт сохраняется по val loss, а log.json переписывается при каждой оценке, так что график можно строить на ходу.

На тестовом прогоне первая строка выглядит примерно так: шаг 0 | train 6.487 | val 6.495 | ... — снова около $\ln V$ (у маленького тестового корпуса словарь получился 651 токен, $\ln 651 \approx 6.48$ ). Дальше loss должен уверенно падать. Скорость в токенах в секунду учитывает и время на оценку, так что TFLOPS получатся слегка заниженными.

Переход между конфигурациями — это замена строки cfg = Config(...) и batch_size. Для S на CPU начните с max_steps = 3000, посмотрите на скорость и пересчитайте, сколько шагов успеете за ночь.

Наш train.py рассчитан на одну карту, и для проекта этого хватает: модель M «по Chinchilla» (517M токенов) при условных 20 TFLOPS обучается около часа, а реальную скорость покажет строка лога. Четыре карты удобнее всего занять четырьмя опытами одновременно — это ровно то, что нужно для уровня Strong. Подготовьте папки run0…run3: в каждой своя копия train.py с нужной конфигурацией, а model.py, train.bin, val.bin и meta.json — ссылки на общие файлы (ln -s ../model.py . и так далее). Затем:

for i in 0 1 2 3; do
  (cd run$i && CUDA_VISIBLE_DEVICES=$i python train.py > train.log 2>&1) &
done; wait

BF16 карты серии RTX 40 поддерживают, train.py включит его сам.
Сравнивайте опыты при равном бюджете FLOPs: произведение $6 \cdot N \cdot \text{токенов}$ у них должно совпадать.

Одна модель на четырёх картах — хорошее упражнение к уроку 5.4.3. Отличия train_ddp.py от train.py:

# запуск: torchrun --nproc_per_node 4 train_ddp.py
import contextlib, os
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
rank, world = dist.get_rank(), dist.get_world_size()
local_rank = int(os.environ["LOCAL_RANK"])
device = f"cuda:{local_rank}"
torch.cuda.set_device(device)
torch.manual_seed(42 + rank)          # свои окна на каждой карте, иначе все четыре учатся на одних и тех же
model = DDP(TinyLLM(cfg).to(device), device_ids=[local_rank])    # веса DDP скопирует с rank 0
use_bf16 = torch.cuda.is_bf16_supported()   # старая проверка device == "cuda" для "cuda:1" дала бы False
tokens_per_step = batch_size * accum * cfg.max_seq_len * world
# в цикле накопления — model.no_sync() на всех микробатчах, кроме последнего (урок 5.4.3);
# evaluate() — через model.module; print, log.json и torch.save(model.module.state_dict()) — только при rank == 0

Сверьте токены в секунду с запуском на одной карте: это ваша эффективность масштабирования на PCIe.

Большой прогон на ночь. По таблице из урока 5.1.1 четыре карты за ночь дают порядка $3 \cdot 10^{18}$ FLOPs. Попробуйте конфигурацию L: dim=768, n_layers=12, n_heads=12, n_kv_heads=4, ffn_dim=2048, это 80.4M параметров. По Chinchilla ей нужно около 1.6B токенов, то есть $6 \cdot 80.4 \cdot 10^6 \cdot 1.6 \cdot 10^9 \approx 7.7 \cdot 10^{17}$ FLOPs: от 2 до 4 часов на четырёх картах в зависимости от вашего MFU (оценка). Узкое место здесь — данные, а не карты: нужны миллионы документов FineWeb-2, а не сотни тысяч. Сколько гигабайт текста это значит, посчитайте по bytes_per_token из meta.json, а токенизацию раздайте по ядрам EPYC. Настройка карт — на странице «Ваша лаборатория (4×GPU)».

Шаг 5. Базовые линии, графики и генерация

Цифра «val loss = 3.9» сама по себе ничего не говорит. Сравните её с моделями, которые точно не умеют пользоваться длинным контекстом: униграммной (частоты токенов) и биграммной (частоты пар). Это ваши базовые линии:

import json, math
import numpy as np

V = json.load(open("meta.json"))["vocab_size"]
train = np.fromfile("train.bin", dtype=np.uint16).astype(np.int64)
val = np.fromfile("val.bin", dtype=np.uint16).astype(np.int64)

# Униграммы: частоты токенов (+1 к каждому, чтобы не было нулей)
uni = np.bincount(train, minlength=V) + 1.0
p_uni = uni / uni.sum()
loss_uni = -np.log(p_uni[val]).mean()

# Биграммы со сглаживанием Дирихле: p(w | v) = (c(v, w) + mu · p_uni(w)) / (c(v) + mu).
# Это родственник сглаживания +α из урока 4.1.2: к счётчику пары подмешиваем частоту униграммы, mu — сила подмешивания.
mu = 100.0
# Трюк: пару (v, w) кодируем одним числом v·V + w — как номер клетки в таблице V×V (урок 4.1.2).
# np.unique возвращает отсортированные уникальные коды пар и сколько раз каждая встретилась
pairs, counts = np.unique(train[:-1] * V + train[1:], return_counts=True)
c_prev = np.bincount(train[:-1], minlength=V)           # c(v): сколько раз токен стоял «перед» чем-то
keys = val[:-1] * V + val[1:]                           # коды пар валидации
# searchsorted ищет место кода в отсортированном pairs; если код там действительно есть — берём счётчик, иначе 0
pos = np.clip(np.searchsorted(pairs, keys), 0, len(pairs) - 1)
c_pair = np.where(pairs[pos] == keys, counts[pos], 0)
p_bi = (c_pair + mu * p_uni[val[1:]]) / (c_prev[val[:-1]] + mu)
loss_bi = -np.log(p_bi).mean()

print(f"ln V      = {math.log(V):.3f}  (модель, которая ничего не знает)")
print(f"униграммы = {loss_uni:.3f}")
print(f"биграммы  = {loss_bi:.3f}")

Трансформер должен быть ниже обеих линий. Если он застрял на уровне униграмм — модель выучила только частоты слов и не смотрит на контекст: ищите ошибку в маске внимания, в сдвиге меток или в learning rate.

График loss по log.json:

import json
import matplotlib.pyplot as plt

log = json.load(open("log.json"))
tokens = [r["tokens"] / 1e6 for r in log]
plt.plot(tokens, [r["train"] for r in log], label="train")
plt.plot(tokens, [r["val"] for r in log], label="val")
plt.xlabel("токенов, млн")
plt.ylabel("loss, нат/токен")
plt.legend()
plt.grid(alpha=0.3)
plt.savefig("loss.png", dpi=120, bbox_inches="tight")
print("сохранено: loss.png")

И генерация из лучшего чекпоинта:

import torch
from tokenizers import Tokenizer
from model import Config, TinyLLM

ckpt = torch.load("ckpt.pt", map_location="cpu")
model = TinyLLM(Config(**ckpt["cfg"]))
model.load_state_dict(ckpt["model"])
model.eval()
tok = Tokenizer.from_file("tokenizer.json")

torch.manual_seed(0)
for prompt in ["Мудрая сова", "Мальчик Ван"]:
    ids = torch.tensor([tok.encode(prompt).ids])
    out = model.generate(ids, max_new_tokens=40, temperature=0.8, top_k=50)
    # skip_special_tokens=False: иначе decode молча выбросит <|endoftext|>, и обрезать по нему будет нечего
    print(tok.decode(out[0].tolist(), skip_special_tokens=False).split("<|endoftext|>")[0])

Бит на байт. train.py печатает ещё одну метрику — бит на байт:

\mathrm{BPB} = \frac{\text{loss (нат/токен)}}{\ln 2 \cdot \text{байт на токен}}

Loss на токен зависит от токенизатора: с большим словарём токены длиннее, и loss на токен выше при той же модели. BPB пересчитывает всё на байт исходного текста, поэтому так можно сравнивать модели с разными словарями, в том числе вашу и MiniMind. Именно эту метрику советует автор MiniMind.

Пример с числами: loss 3.0 нат/токен при 3 байтах текста на токен даёт $3.0 / (0.693 \cdot 3) \approx 1.44$ бита на байт. У посимвольной модели на русском один токен — это одна буква, то есть 2 байта, поэтому её BPB почти вдвое меньше, чем loss в битах на символ («почти» — потому что пробелы, цифры и знаки препинания занимают по 1 байту).

Разминка в браузере: базовые линии и BPB

Прежде чем запускать baselines.py на своём корпусе, соберите те же формулы на чистом Python и проверьте их на крошечном посимвольном примере «мама мыла раму». Это те же функции, что понадобятся в домашней работе в конце урока.

Базовые линии и бит на байт7 тестов

import math
from collections import Counter


def initial_loss(vocab_size):
    """Loss только что инициализированной модели со словарём vocab_size
    (она предсказывает почти равномерное распределение)."""
    # TODO
    ...


def unigram_probs(train, vocab_size):
    """Список длины vocab_size: p(w) = (c(w) + 1) / (N + V),
    где c(w) — сколько раз токен w встречается в train, N = len(train), V = vocab_size.
    Подсказка: Counter(train)[w] даёт c(w), в том числе 0 для невстреченных токенов."""
    # TODO
    ...


def unigram_loss(train, val, vocab_size):
    """Средний -ln p(w) по всем токенам val, где p — из unigram_probs(train, vocab_size)."""
    # TODO
    ...


def bigram_loss(train, val, vocab_size, mu=100.0):
    """Средний -ln p(w | v) по всем соседним парам (v, w) из val (их len(val) - 1).
    p(w | v) = (c(v, w) + mu · p_uni(w)) / (c(v) + mu), где
      c(v, w) — сколько раз пара (v, w) стоит подряд в train,
      c(v)    — сколько раз v встречается в train[:-1] (то есть стоит перед каким-то токеном),
      p_uni   — unigram_probs(train, vocab_size)."""
    # TODO
    ...


def bits_per_byte(loss_nats, n_tokens, n_bytes):
    """BPB: loss в натах на токен → биты на байт исходного текста.
    Всего натов: loss_nats · n_tokens; в битах — делим на ln 2; на байт — делим на n_bytes."""
    # TODO
    ...


# Проверьте себя кнопкой «Запустить»: посимвольная «модель» на коротком русском тексте
text = "мама мыла раму. рама мыла маму. " * 30
vocab = sorted(set(text))
ids = [vocab.index(ch) for ch in text]
train, val = ids[:800], ids[800:]
V = len(vocab)
print("V =", V, "ln V =", initial_loss(V))
print("униграммы:", unigram_loss(train, val, V))
print("биграммы: ", bigram_loss(train, val, V, mu=1.0))

На этом примере биграммы обгоняют униграммы больше чем втрое (0.56 против 1.87 нат/символ): в тексте из повторяющихся слов следующая буква почти всегда определяется предыдущей. На настоящем корпусе разрыв меньше, а трансформер должен уйти ниже обеих линий.

Как читать кривые

Что видите	Вероятная причина	Что делать
Начальный loss намного больше $\ln V$	Плохая инициализация или ошибка в сдвиге меток	Проверьте `std = 0.02` и то, что `y` сдвинут на 1 относительно `x`
Loss стоит на месте с первых шагов	Слишком маленький LR или градиенты не доходят	Проверьте `opt.step()`, `zero_grad`, LR после разогрева
Loss взлетает или становится `nan`	Слишком большой LR, нет обрезки градиента	Уменьшите `lr_max` в 2–3 раза, проверьте `clip_grad_norm_`
train падает, val растёт	Переобучение: мало данных, много эпох	Больше данных, меньше шагов, dropout 0.1
val заметно ниже train	Утечка (дубликаты между train и val) или слишком маленькая и «лёгкая» val-выборка	Дедупликация, разбиение по документам, более крупная val-выборка
Генерация повторяет одну фразу	Модель мала или недообучена, слишком низкая температура	Учите дольше, поднимите температуру до 0.9–1.0

Дополнительно: MoE-версия

Для уровня Boss попробуйте заменить FFN на смесь экспертов с 4 экспертами и маршрутизацией top-1, как в minimind-3-moe. Код подключается к модели выше без изменений в остальных классах:

import torch
import torch.nn as nn
import torch.nn.functional as F
from model import Config, SwiGLU, TinyLLM

class MoE(nn.Module):
    """n_experts экспертов SwiGLU, маршрутизация top-1 (как в minimind-3-moe)."""
    def __init__(self, cfg, n_experts=4):
        super().__init__()
        self.router = nn.Linear(cfg.dim, n_experts, bias=False)
        self.experts = nn.ModuleList(SwiGLU(cfg) for _ in range(n_experts))
        self.aux_loss = torch.tensor(0.0)

    def forward(self, x):
        B, T, D = x.shape
        flat = x.reshape(-1, D)
        probs = F.softmax(self.router(flat), dim=-1)          # (B·T, E)
        weight, idx = probs.max(dim=-1)                       # top-1: лучший эксперт и его вероятность
        out = torch.zeros_like(flat)
        for e, expert in enumerate(self.experts):
            sel = idx == e
            if sel.any():
                out[sel] = expert(flat[sel]) * weight[sel, None]
        # aux loss из Switch Transformer: E · Σ (доля токенов у эксперта) · (средняя вероятность эксперта)
        frac = F.one_hot(idx, len(self.experts)).float().mean(0)
        self.aux_loss = len(self.experts) * (frac * probs.mean(0)).sum()
        return out.view(B, T, D)

cfg = Config(vocab_size=6400, dim=384, n_layers=6, n_heads=6, n_kv_heads=2, ffn_dim=1024)   # S
model = TinyLLM(cfg)
for blk in model.blocks:
    blk.ffn = MoE(cfg)
model.apply(model._init)        # новые эксперты и роутер — с той же инициализацией std = 0.02
# в цикле обучения:  loss = loss + 0.01 * sum(b.ffn.aux_loss for b in model.blocks) / len(model.blocks)

У конфигурации S с MoE 33.1M параметров всего и те же ~11.9M активных. При идеальном балансе aux_loss равен 1, при коллапсе на одного эксперта — 4. Честное сравнение с плотной моделью — при равных активных параметрах и равном числе токенов: тогда FLOPs обучения почти одинаковы (урок 5.1.1). Ожидайте, что в чистом PyTorch MoE будет обучаться медленнее, — та же картина, что у автора MiniMind.

Ориентиры уровней

Проект оценивается ступенями по образцу курсов Ли Хунъи. Таблица ниже — ориентиры для самопроверки; точные проверки и пороги домашней работы описаны в её README.

Уровень	Что сделать	Признак успеха
Simple	Конвейер работает от `corpus.jsonl` до генерации на XS	Начальный loss около $\ln V$ ; модель переобучается на одном батче почти до нуля; `sample.py` печатает текст
Medium	Полное обучение конфигурации S на CPU на своём русском корпусе	val loss ниже униграммной линии; график loss без всплесков; в генерации настоящие русские слова
Strong	Аккуратный эксперимент	val loss ниже биграммной линии; посчитан BPB; сравнены две конфигурации при равном бюджете FLOPs (например, глубокая и широкая или GQA и MHA)
Boss	Обучение на GPU	Конфигурация M на сотнях миллионов токенов (около оптимума Chinchilla) или MoE против плотной модели при равных активных параметрах; метрики сохранены в `metrics.json`

Китайский угол: «ляньдань» — алхимия обучения моделей

Слово «ляньдань» («выплавлять пилюлю бессмертия», «алхимия») и «ляньдань-ши» («мастер алхимии») вы встречали в уроке 3.3.3: так в Китае шутливо называют обучение нейросетей и подбор гиперпараметров. Теперь вы сами у «печи». Datawhale так и описывает свою китайскую адаптацию Stanford CS336 diy-llm — «мастерская LLM-алхимии». Шутка точная: пока у вас нет графиков, базовых линий и чекпоинтов, обучение похоже на гадание.

Ваш проект повторяет путь двух китайских учебных проектов:

MiniMind — те же решения: byte-level BPE на 6400 токенов, RMSNorm, RoPE, SwiGLU, GQA и косинусное расписание. Различия — упаковка вместо паддинга и свой русский токенизатор.
happy-llm, глава 5 (github.com/datawhalechina/happy-llm) — сборка LLaMA2-подобной модели на 215M параметров с токенизатором и предобучением, веса выложены на ModelScope. Хорошо читать параллельно, если хотите второй взгляд на те же шаги.

После проекта вы сможете открыть код MiniMind и узнавать в нём каждую строку. Это и есть цель.

Проверь себя

Проверь себя5 вопросов

1 Только что инициализированная модель со словарём 4096 токенов. Какой примерно loss вы ожидаете на первом шаге? Округлите до одного знака.

2 val loss вашей модели равен loss униграммной базовой линии и не падает. Что это значит?Всё хорошо: униграммы — сильная модельМодель выучила только частоты токенов и не использует контекст — вероятна ошибка в маске, сдвиге меток или LRНужно увеличить словарь

3 Какие утверждения про упаковку последовательностей верны?Можно выбрать несколько ответовДокументы склеиваются в один поток через разделитель <|endoftext|>Нет позиций-паддингов, все вычисления идут на настоящие токеныОкно обучения может начаться с середины документаУпаковка обязательна для работы RoPE

4 Loss модели — 2.0 нат/токен, в среднем на токен приходится 4 байта текста. Чему равен BPB? Округлите до двух знаков.

5 train loss падает, а val loss после середины обучения начал расти. Что вероятнее всего?Утечка тестовых данныхПереобучение: модель запоминает обучающие данные — мало данных или много эпохСлишком маленький learning rateОшибка в токенизаторе

Итоги

Проект — тот же конвейер, что у больших моделей: данные → токенизатор → модель → обучение → оценка, только в масштабе одной машины.
Сначала бюджет: $C = 6ND$ и ваши реальные TFLOPS показывают, какую модель и сколько токенов вы можете себе позволить.
Свой byte-level BPE на русских данных, упаковка вместо паддинга, разбиение train/val по документам.
Модель на 4–26M параметров с RMSNorm, RoPE, SwiGLU и GQA; проверки здравого смысла — loss около $\ln V$ и переобучение на одном батче.
Оценка — не «на глаз»: базовые линии (униграммы, биграммы), BPB и честные кривые train и val.

s5-hw02-tiny-llm-pretrain — всё, что окружает модель в этом проекте, под автотестами. Модель с современным блоком уже готова, а вы пишете данные (разбиение train/val по документам, посимвольный токенизатор, упаковка, get_batch, униграммная и биграммная линии) и цикл обучения (разогрев и косинус, AdamW с группами weight decay, обрезка и накопление градиентов, оценка, чекпоинты, BPB). Корпус — русские тексты уроков курса. Уровни Simple–Strong проверяют код, Boss — полный прогон configs/cpu.json на CPU. Условие, заготовки и тесты — в папке homework/s5-hw02-tiny-llm-pretrain/. Проверка из корня репозитория:

python homework/check.py s5-hw02-tiny-llm-pretrain

Дальше: модуль 5.7 «GPU изнутри» — урок Как устроена видеокарта и куда уходит время. Разберём, во что упирается обучение вашей модели на видеокарте, научимся профилировать шаг, писать свои ядра и честно масштабировать обучение на четыре карты. После него — ступень 6: научим базовую модель отвечать на вопросы (SFT) и рассуждать (RL).