llm-course-hw1

Домашняя работа по обучению byte-level BPE токенизатора и небольшой causal language model на русских анекдотах.

Dataset

Использовался датасет IgorVolochay/russian_jokes, файл dataset.txt.

После разбиения:

split	rows
train	135497
test	15056

Tokenizer

Реализован byte-level BPE tokenizer:

• base vocabulary: 256 byte-level токенов;
• итоговый vocab size: 1024;
• special token: [EOS].

Model

Основная модель: decoder-only Transformer для causal language modeling.

Использованные компоненты:

• RMSNorm;
• causal self-attention;
• Grouped-Query Attention;
• ALiBi positional bias;
• SwiGLU feed-forward block;
• AdamW;
• linear warmup + linear decay scheduler.

Конфигурации:

model	layers	heads	kv heads	hidden dim	intermediate dim
nano	3	4	2	96	256
mini	6	6	3	384	1024
small	12	12	6	768	2048

Hyperparameter Search

Короткий поиск по mini и small:

model	class	steps	batch	lr	wd	dropout	train_loss	val_loss
mini	TransformerForCausalLM	700	24	2.0e-04	0.01	0.10	4.4266	4.3444
mini	TransformerForCausalLM	700	24	3.0e-04	0.01	0.10	4.1954	4.1069
mini	TransformerForCausalLM	700	24	5.0e-04	0.05	0.15	4.1432	4.0272
small	TransformerForCausalLM	700	12	2.0e-04	0.01	0.10	4.0482	4.1350
small	TransformerForCausalLM	700	12	3.0e-04	0.01	0.10	4.0802	4.0462
small	TransformerForCausalLM	700	12	5.0e-04	0.05	0.15	4.2554	4.3985

Кратко:

• mini лучше всего сработал при более агрессивном режиме lr=5e-4, weight_decay=0.05, dropout=0.15.
• small хуже перенес агрессивный lr=5e-4; лучший короткий результат был при lr=3e-4, weight_decay=0.01, dropout=0.10.

First Final Training

model	steps	batch	lr	wd	dropout	train_loss	val_loss
mini	6000	24	5.0e-04	0.05	0.15	3.1870	3.0259
small	8000	12	3.0e-04	0.01	0.10	2.8207	2.8535

Refined Hyperparameter Search

model	class	steps	batch	lr	wd	dropout	train_loss	val_loss
mini	TransformerForCausalLM	3000	24	4.0e-04	0.03	0.12	3.4369	3.2428
mini	TransformerForCausalLM	3000	24	5.0e-04	0.03	0.12	3.3032	3.2085
mini	TransformerForCausalLM	3000	24	6.0e-04	0.05	0.15	3.2975	3.2324
small	TransformerForCausalLM	3000	12	2.5e-04	0.01	0.10	3.3234	3.2166
small	TransformerForCausalLM	3000	12	3.0e-04	0.01	0.12	3.4603	3.2327
small	TransformerForCausalLM	3000	12	3.5e-04	0.02	0.12	3.2497	3.2388

Лучшие refined-настройки:

{
  "mini": {
    "learning_rate": 0.0005,
    "weight_decay": 0.03,
    "dropout": 0.12
  },
  "small": {
    "learning_rate": 0.00025,
    "weight_decay": 0.01,
    "dropout": 0.1
  }
}

Deep Final Training

model	steps	batch	lr	wd	dropout	train_loss	val_loss
mini	10000	24	5.0e-04	0.03	0.12	2.9914	2.8547
small	12000	12	2.5e-04	0.01	0.10	2.7788	2.7377

Лучший итоговый checkpoint по validation loss: small (val_loss=2.7377).

Generation Examples

mini

Prompt: Заходит в бар

Заходит в бар, видит. Потом сидит на кухне. Дедушка снимит, и всякий, пьяный и говорит: - Помимся, как я тебе говори! Другой: - Папа, я тебе папа, я пьянка на кухню. Сталиной, погоди, что ты не пришла в гости и выражает. - Да, но ты не знаешь, я не знаю. - Дедушка, это я, сынок! - Опять!

Prompt: Штирлиц пришел домой

Штирлиц пришел домой с работы, в ней котов и говорит: - Папа, я не хочу, чтобы ты согласна в школу! А я и с ним, что я пойму в кафе...

Prompt: Студент на экзамене

Студент на экзамене:- У тебя что-то вечером встречаются с денег.- Они не привлекают.

small

Prompt: Заходит в бар

Заходит в бар. Компания, смотрит и говорит: "Абрам, как это за рубашку в кинуле!" Все в порядке воскресенье, наверно, весь вниз, достает рабочий ванной в сторону. Сидят двое и толпатных, встречает одна и говорит: - Вовочка, ты кем? - Я, блин, куда?

Prompt: Штирлиц пришел домой

Штирлиц пришел домой. Жена:- Слушай, а ты знаешь, что я таки даже не хочу.Муж:- Ты что, дорогая, ты не волнуйся, я не могу.

Prompt: Студент на экзамене

Студент на экзамене:- Сколько у вас в школе?- Да.- Ага, как я вышла!

Decoding Search

mini

temperature=0.8, top_k=20

Заходит в бар, видит - улыбается на работу, а жена спрашивает: - Мужик, что это вы делаете? - Да вот, не волнуйтесь, какой ты здесь волосы плохо!

temperature=0.9, top_k=30

Заходит в бар и кричит: - Сара, а зачем тебе так устроился? Мимо доказывает: - Тогда, кем ты был на работу? - Ну, пусть же я плачу.

temperature=0.95, top_k=40

Заходит в бар, смотрит, значит, по пустым метро стоит за окном. Мы ему лжично раздражать в туалете - и шо, у вас там, наверное, дорогая, ну на хлеб поездце. Представляешь, что я такое в этой камере. Там некогда нет. Бедно превращается, а я рыба - чушью и кладут с ней уже в кварт

small

temperature=0.8, top_k=20

Заходит в бар с барменом, а там в баре с папой падают. - Чего, сюда едешь? - Да я знаю, что ты на почту упала.

temperature=0.9, top_k=30

Заходит в бар. Подходит маленькая девушка:- Товарищ мерседес...Девушка:- У вас есть базар, пиво?- Нет, солнцы!

temperature=0.95, top_k=40

Заходит в бар к своему суду и говорит: - Ну сколько будет? - Да поймать. - ? - Да, да. - А вы?

Bonus Experiments

Дополнительно реализованы RoPE и учебная версия MLA с low-rank KV compression.

model	class	steps	batch	lr	wd	dropout	train_loss	val_loss
nano	TransformerForCausalLMRoPE	700	24	3.0e-04	0.01	0.10	5.1198	5.0737
nano	TransformerForCausalLMMLA	700	24	3.0e-04	0.01	0.10	5.0007	4.9573

Loaded Check Example

Prompt: Штирлиц пришел домой

Штирлиц пришел домой в командировку. Жена говорит мужу:- Спрашиваете как-нибудь веселое.- Отличное доступное, что ли?- Классно. Сегодня я выбил, а он ведь не мог бы выбрать.

References

• RoFormer: Enhanced Transformer with Rotary Position Embedding
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
• Hugging Face Hub upload guide