Всё началось с наивной мысли: зачем платить за API или тащить 7B-модель, если мне нужна маленькая модель для простых разговоров на одном языке? Логика казалась железной — большие модели умеют всё и на всех языках сразу, но это же избыточно. 0.7B, заточенная под один язык и один стиль общения, должна справляться не хуже.

Спойлер: это было наивно. Но путь оказался ценнее результата.

В этой статье — как я прошёл путь от стандартного nanoGPT до кастомной архитектуры с RoPE/SwiGLU/GQA, собрал русскоязычный корпус с нуля, и придумал распределённое обучение на бесплатных Colab-воркерах через Google Drive.

Почему не взять готовую модель?

Честный ответ: потому что хотел разобраться как оно работает изнутри. Взять Qwen или Llama, запустить файн-тюнинг — это понятно и работает. Но когда ты сам строишь модель с нуля, каждая архитектурная деталь перестаёт быть магией.

Плюс была конкретная задача: нужен персонаж с определённым стилем речи, на русском языке. Казалось логичным — меньше модель, меньше ресурсов, проще управлять поведением. Оказалось, что 0.7B это очень мало для нормальной генерации связного текста, но это я понял позже.

Датасет: какой язык похож на нужный мне

Перед сбором данных встал неочевидный вопрос: какой стиль русского языка нужен? Официальные новости — нет. Научные тексты — нет. Нужен был живой, разговорный, эмоциональный язык.

Корпус собрал из трёх частей:

Taiga — готовый русскоязычный корпус: новости, журналы, художественная литература, субтитры. Хорошая база, но стиль неоднородный.

Собственный скрейп — ~566k документов игровых медиа, блог-платформ, литературных сообществ. Именно здесь живой язык: эмоции, сленг, неформальные обсуждения.

FineWeb2 (rus_Cyrl) — веб-корпус, но сырой. Его пришлось фильтровать стриминговым пайплайном.

Фильтрация FineWeb2

Веб-данные — это мусор по умолчанию. SEO-тексты, прайс-листы, резюме, битые символы. Написал стриминговый фильтр чтобы не грузить всё в память:

@dataclass
class Config:
    min_chars: int = 50
    max_chars: int = 5000
    min_letter_ratio: float = 0.4      # минимум 40% букв в тексте
    max_special_ratio: float = 0.25    # не больше 25% спецсимволов
    lang_conf_threshold: float = 0.6   # уверенность langid
    simhash_threshold: int = 3         # порог для дедупликации

Плюс список стоп-слов для быстрой фильтрации (“резюме”, “прайс-лист”, “seo”, “вакансия” и т.д.) — дешевле langid и отсекает большую часть мусора до тяжёлых проверок.

Итого: ~12B токенов после фильтрации.

Токенизатор

Кастомный BPE на 51 200 токенов, обученный на русскоязычном корпусе. GPT-2-шный токенизатор плохо работает с русским — слова разбиваются на слишком мелкие куски, контекстное окно расходуется неэффективно.

Эволюция архитектуры: пять итераций

Итерация 1 — GPT-2 Small, базовый старт

Первая модель — чистый nanoGPT без изменений. 124M параметров, fp32, без компиляции.

batch_size=8, block_size=1024, grad_accum=1, iters=600k
train loss: 3.99 | val loss: 4.02

Работает, обучается, loss падает. Но медленно и генерация слабая — слишком мало данных. Следующий шаг — добавил torch.compile и переход на fp16. Только это дало заметный прирост скорости без изменения архитектуры.

Итерация 2 — GPT-2 Small, динамические гиперпараметры

Главный эксперимент: менять гиперпараметры прямо в процессе обучения, не перезапуская. Постепенно увеличивал batch size и gradient accumulation, подбирал weight decay.

block_size=1024
batch size: 8 → 16 → 20
grad accum: 3 → 6 → 9 → 12 → 24
weight decay: 0.1 → 0.01 → 0.05 → 0.1
токенов: ~5.3B
train loss: 3.24 | val loss: 3.30

Существенный прогресс. Большой gradient accumulation = большой эффективный батч = стабильнее обучение. Но удар по производительности ощутимый.

Итерация 3 — GPT-2 Medium, переход на 345M

Та же архитектура, больше параметров. Добавил dropout во второй половине обучения.

batch_size=4, grad_accum: 20 → 60
dropout: 0 → 0.05 (со второй половины)
токенов: ~5.5B
train loss: 3.07 | val loss: 3.12

Итерация 4 — GPT-2 Large, распределённое обучение

GPT-2 Large (774M) технически влезал в одну Colab-сессию, но только с gradient checkpointing. Это когда активации не сохраняются при forward и пересчитываются заново при backward. Память экономится, но скорость падает примерно в 2 раза. На практике обучение стало невыносимо медленным.

Именно из-за скорости появилась распределённая схема — об этом отдельный раздел ниже.

n_layer=36, n_head=20, n_embd=1280
block_size=2048, batch_size=2, grad_accum=125
токенов: ~15B
train loss: 2.86 | val loss: 3.04

При 5 воркерах ускорение вышло примерно x2 относительно одного — казалось бы немного, но на практике это разница между “обучение идёт” и “обучение стоит”. 15B токенов и 774M параметров — уже чувствуется в качестве генерации.

Итерация 5 — кастомная архитектура

Финальная модель — переработанная архитектура с современными компонентами. По сути, это то, чем отличается современный LLaMA-стиль от оригинального GPT-2.

токенов: ~5.1B
train loss: 3.07 | val loss: 3.11

Loss похож на итерацию 3, но модель меньше токенов видела — архитектура эффективнее использует параметры. И инференс заметно быстрее.

Детали архитектуры

RoPE

Вместо таблицы абсолютных позиций — поворот векторов Q и K в зависимости от позиции. Относительное расстояние между токенами кодируется в самом attention.

def apply_rotary_pos_emb(q, k, cos, sin):
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

SwiGLU

Вместо одной матрицы в FFN — две, с gate-механизмом:

def forward(self, x):
    return self.w3(self.swish(self.w1(x)) * self.w2(x))

hidden_dim = int(4 * d_model * 2/3) — стандартная формула чтобы компенсировать дополнительную матрицу и сохранить примерно тот же FLOPs.

GQA

Grouped Query Attention: вместо отдельных KV-голов для каждой Q-головы — одна группа KV на несколько Q. Меньше памяти, быстрее — особенно при инференсе с длинным контекстом.

if self.n_kv_head != self.n_head:
    k = k.repeat_interleave(self.n_head // self.n_kv_head, dim=1)
    v = v.repeat_interleave(self.n_head // self.n_kv_head, dim=1)

Selective Gradient Checkpointing

Полный gradient checkpointing режет память в 2x, но замедляет обучение на ~30% — пересчитываются активации при backward pass. Решение: чекпоинтить только часть слоёв.

checkpoint_strategy: str = "custom"  # "alternate", "first_last", "custom", "all"

На практике чекпоинтинг чётных слоёв давал хороший баланс памяти и скорости.

Распределённое обучение на Colab через Google Drive

Это самая нестандартная часть проекта. GPT-2 Large с полным gradient checkpointing обучался слишком медленно на одной сессии. Нужно было распараллелить.

Стандартный DDP через NCCL или Gloo требует прямой связи между нодами. В Colab это невозможно — у сессий нет постоянных IP, нет общей сети. Нужен был другой транспорт.

Идея: Google Drive как шина градиентов

Google Диск в Colab монтируется как обычная файловая система. Каждый воркер может писать и читать файлы. Значит, можно передавать градиенты через файлы.

Схема:

Воркер 1                          Воркер 2
─────────                         ─────────
forward pass                      forward pass
backward pass                     backward pass
сохранить grad_iter_N_worker_1.pt сохранить grad_iter_N_worker_2.pt
       ↓                                 ↓
       ожидание файла воркера 2          ожидание файла воркера 1
       ↓                                 ↓
загрузить, усреднить, применить  загрузить, усреднить, применить
optimizer.step()                  optimizer.step()

Каждый воркер знает своё worker_id и общее count_workers. После backward pass — сохраняет градиенты под своим ID и номером итерации, затем ждёт файлы остальных.

# Сохранение своих градиентов
save_gradients(model, f"drive/.../grad_iter_{local_iter_num}_worker_{worker_id}.pt")

# Ожидание и загрузка чужих
paths = [f"grad_iter_{local_iter_num}_worker_{i}.pt"
         for i in range(1, count_workers + 1) if i != worker_id]
accumulate_gradients_into_model(model, paths)

Усреднение делает каждый воркер самостоятельно:

def accumulate_gradients_into_model(model, gradient_paths):
    for path in gradient_paths:
        grads = torch.load(path, map_location='cpu')
        for param, g in zip(model.parameters(), grads):
            if param.grad is None:
                param.grad = g.clone()
            else:
                param.grad.add_(g)
    # Усреднение по всем воркерам
    for param in model.parameters():
        if param.grad is not None:
            param.grad.div_(len(gradient_paths) + 1)  # +1 свои

Инфраструктура: несколько аккаунтов Google

Каждый воркер — это отдельная Colab-сессия под отдельным Google-аккаунтом. Общая точка синхронизации — папка на Google Drive, открытая для всех аккаунтов. Каждый воркер монтирует свой Drive, но у всех есть ярлык на эту общую папку — туда и пишутся градиенты.

Два процесса на каждом воркере

Внутри одной Colab-сессии работало два процесса. Обучение запускалось через subprocess.Popen — фоново, не блокируя ячейку:

import subprocess
proc = subprocess.Popen(["python", "train.py"])

Вторая ячейка крутила даунлоадер — следила за появлением файлов градиентов других воркеров на Google Drive и скачивала их локально.

Флаги как сигнал готовности

Важная деталь про Google Drive: файлы на диске всегда целые — незавершённая загрузка туда просто не синхронизируется. Но при скачивании обратно локально файл может появиться раньше, чем докачается полностью.

Поэтому даунлоадер, завершив скачивание файла, создавал локальный .flag:

# даунлоадер: скачали файл → создаём флаг
with open(f"{local_iter_num}.flag", "w") as f:
    pass

Обучение ждало флага, а не самого файла с градиентами:

# обучение: ждём флага от даунлоадера
while not os.path.exists(f"{local_iter_num}.flag"):
    time.sleep(0.1)

Флаг — гарантия что файл градиентов уже полностью лежит локально и можно читать.

Градиенты в float16

Файлы градиентов для Large модели весят несколько сотен мегабайт. Сохранять в float32 — долго грузить и долго синхронизировать. Решение: сохранять в float16, конвертировать обратно при загрузке.

grads = [param.grad.detach().cpu().to(torch.float16) if param.grad is not None else None
         for param in model.parameters()]
torch.save(grads, path)

Ограничения подхода

• Скорость синхронизации: ожидание Drive — самый медленный шаг. Несколько секунд на итерацию накапливаются.

• КПД распараллеливания: 5 воркеров дали x2, а не x5. Большую часть времени воркеры просто ждут друг друга.

• Нестабильность: Colab-сессии падают без предупреждения. Обязательно уметь возобновлять с чекпоинта.

• Масштаб: больше 5-6 воркеров — координация через файлы становится узким местом.

Для прода это не годится. Для бесплатного обучения на Colab — работает.

Результаты и выводы

Что я понял

0.7B — это мало. Для нормальных связных диалогов нужно минимум 3-7B, и это при условии хорошего файн-тюнинга поверх. Идея что “маленькая специализированная модель заменит большую” работает хуже чем хотелось бы.

Данные важнее архитектуры. v4 с простой GPT-2 Large архитектурой и 15B токенов показала лучший loss, чем v5 с современной архитектурой и 5B токенов. Качество корпуса и объём данных бьют архитектурные улучшения на небольших масштабах.

Распределённое обучение — не только про скорость. Два воркера с синхронизацией через Drive дали возможность обучать модель, которая физически не помещалась в одну сессию. Это важнее ускорения.

Chinchilla-оптимум достигли только v2 и v4. ~20 токенов на параметр — и это объясняет их отрыв от остальных итераций лучше, чем любые архитектурные детали. Остальные просто недообучены по данным.