--- language: ru tags: - causal-lm - byte-level-bpe - transformer - gqa - alibi - swiglu - rope - mla library_name: pytorch --- # llm-course Проект по обучению byte-level BPE токенизатора и небольшой causal language model на русских анекдотах. ## Dataset Использовался датасет [`IgorVolochay/russian_jokes`](https://huggingface.co/datasets/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-настройки: ```json { "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](https://arxiv.org/abs/2104.09864) - [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://arxiv.org/abs/2405.04434) - [Hugging Face Hub upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload)