maf1oz1/llm-course-hw1

Russian Jokes LM — mini (ALiBi + GQA + SwiGLU)

Маленькая языковая модель, обученная с нуля генерировать русские анекдоты. Домашнее задание №1 курса по большим языковым моделям (HSE): собственный byte-level BPE токенизатор и собственная реализация трансформера-декодера на PyTorch, без готовых библиотек моделей.

В этом репозитории лежат основная модель и токенизатор:

• model.safetensors, config.json — веса и конфиг модели;
• vocabulary.json, merges.json — словарь и слияния BPE-токенизатора.

Токенизатор

• Byte-level BPE (как в GPT-2), реализован и обучен с нуля на обучающей части датасета.
• Размер словаря 1024 токена, включая специальный [EOS].
• Текст предварительно режется на слова регуляркой в стиле GPT-2, слияния применяются внутри слов по приоритету (рангу) обучения.

Архитектура

Трансформер-декодер в стиле LLaMA:

Параметр	Значение
Параметров	10.54M
Слоёв	6
Скрытая размерность	384
Голов внимания (Q / KV)	6 / 3 — Grouped-Query Attention
Позиционное кодирование	ALiBi (без позиционных эмбеддингов)
Feed-forward	SwiGLU, промежуточная размерность 1024
Нормализация	RMSNorm, pre-norm
Контекст	128 токенов
Dropout	0.1

Обучение

• Данные: IgorVolochay/russian_jokes, разбиение 90/10: 135 497 анекдотов в обучении, 15 056 в тесте.
• Задача: предсказание следующего токена, паддинги исключены из лосса по маске.
• Оптимизатор: AdamW, lr 3e-4, weight decay 0.01, gradclip 1.0.
• Расписание lr: warmup первые 10% шагов, затем линейное падение до нуля.
• Батч 16, 25 404 шага (3 эпохи), GPU T4 на Kaggle.

Результаты

Все эксперименты — конфиг mini, одинаковые данные и гиперпараметры (3 эпохи); отличается только механизм внимания и позиционное кодирование.

Модель	Val loss	Perplexity
ALiBi + GQA — эта модель	2.651	14.2
RoPE + GQA (бонус)	2.612	13.6
ALiBi + Multi-Head Latent Attention (бонус)	2.682	14.6

• RoPE немного лучше ALiBi (−0.04 по лоссу): на коротком контексте в 128 токенов разница между способами кодирования позиций небольшая.
• MLA немного хуже (+0.03), но хранит в KV-кэше на токен в каждом слое 128 чисел вместо 384 у GQA — в 3 раза меньше. Реализован упрощённый вариант: K и V восстанавливаются из общего латентного вектора размера 128, без decoupled RoPE и без сжатия запросов; позиции кодируются ALiBi.

Бонусные модели лежат в отдельных репозиториях: maf1oz1/llm-course-hw1-rope и maf1oz1/llm-course-hw1-mla.

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

Без отбора, temperature=0.8, top_k=20, seed=42:

Заходит в бар, значит, в камеру, снимает бутылку водки и спрашивает: — Что случилось? — Сталкивает лекарство, что ли?

Приходит мужик к врачу и говорит: — Я так долго не держу, плохо тороплюсь. — О! Мне уже два дня!

Встречаются два программиста в квартире. Один другому: — Скажите, а вы сегодня в банкете? — Спасибо, не знаю, что это у вас серьезное счастье.

Штирлиц пришел домой, а там его жена ушла, захотела на плохую стормозить. — Спасибо, не надо! — Извини! — А ты что? — А я тебе кажется, что ли!

С параметрами из ячейки проверяющего (temperature=1.0, top_k=10):

Штирлиц пришел домой на рынку. Папа: — Папа, я вчера прошел. Я бы хотел, чтобы я буду ногу! — Ну, и как это?

Модель уверенно воспроизводит форму анекдота — зачины, диалоги через тире, реплики персонажей, завершение через [EOS], — но смысл и юмор пока не держит.

Как запустить

Классы ByteLevelBPETokenizer и TransformerForCausalLM определены в тетрадке домашнего задания — их нужно выполнить перед загрузкой.

import torch

tokenizer = ByteLevelBPETokenizer.from_pretrained("maf1oz1/llm-course-hw1")
model = TransformerForCausalLM.from_pretrained("maf1oz1/llm-course-hw1").eval()

ids = torch.tensor([tokenizer.encode("Штирлиц пришел домой", add_eos_token=False)])
out = model.generate(ids, max_new_tokens=150, eos_token_id=tokenizer.eos_token_id,
                     do_sample=True, top_k=20, temperature=0.8)
print(tokenizer.decode(out[0].tolist()))

Если промпт закодирован с [EOS] на конце (поведение encode по умолчанию), generate отрезает его сам: модель не обучалась продолжать текст после конца анекдота.

Ограничения

• Около 10M параметров и ~10M токенов данных: грамматика в основном правильная, но сюжет быстро теряет смысл.
• Словарь из 1024 токенов режет слова на слоги, поэтому иногда появляются несуществующие слова.
• Анекдоты длиннее 128 токенов при обучении обрезались.
• Модель обучена на анекдотах из интернета и может воспроизводить грубые или стереотипные шутки из датасета.