llm-course-hw1

ДЗ №1. Тут есть byte-level BPE токенизатор и маленькая языковая модель, которые написаны с нуля на PyTorch и обучены на русских анекдотах из датасета IgorVolochay/russian_jokes

Токенизатор

Словарь на 1024 токена: 256 байтов, 767 слияний и служебный [EOS]. Обучался только на train части датасета. В среднем анекдот получается примерно в 73 токена, поэтому контекста в 128 хватает почти всегда. Словарь находится в vocabulary.json, мержи в merges.json

Модель

Decoder-only трансформер в конфигурации mini, 10.52M параметров, 6 слоев и скрытая размерность 384. Используется ALiBi, внимание с GQA (6 голов на запросы и 3 на ключи и значения), в feed-forward блоке SwiGLU с промежуточной размерностью 1024. Перед каждым блоком RMSNorm, dropout 0.1, длина контекста 128 токенов.

Обучение и качество

Из датасета брался dataset.csv, после разбиения с seed 0xC0FFEE вышло 135 497 анекдотов на обучение и 15 056 на валидацию. Оптимизатор AdamW с weight decay 0.01, learning rate первые 10% шагов линейно растет, потом линейно падает до нуля, градиенты обрезаются по норме 1.0

Всего было три запуска. Сначала nano на 0.5M параметров (батч 16, lr 3e-4, 10 000 шагов). Она дошла до validation loss 3.499 и выдавала слова почти без связи. Потом mini с такими же настройками на 20 000 шагов дала 2.693. Лучше всего получилось с батчем 64 и lr 6e-4. В этом случае за 15 000 шагов модель прошла по данным примерно 7 раз и validation loss опустился до 2.487. Эта версия лежит в репозитории. Обучение заняло около 2х часов на T4 в Google Colab, переобучения не было, train и validation loss почти совпадали до самого конца

По текстам можно понять, что модель поняла, как устроен анекдот, т.е там пристутсвуют завязка, диалог через дефисы, концовка и [EOS]. Слова и грамматика в целом нормальные, но смысла и юмора не сильно много (но для 10M параметров это ожидаемо)

Примеры

Генерация с do_sample=True, top_k=10 и температурой 1.0

"Заходит в бар"

• Заходит в бар, а там сидит мужик и думает, что он в баре. Она ему говорит:- Подними штаны, пожалуйста.- А почему ты не в туалет?!- Я не понял, я - адекватный барабан, что ли!

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

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

• Штирлиц пришел домой и поставил в пистолет. Заходит к нему и спрашивает: - Сынок, куда пропадать? - Да вот, папа. - А почему ты не сразу так плохо выкину? - Так в школу, сынок. В семье, восемь лет назад, а в пустыне пятый год.

• Штирлиц пришел домой в песочницу, но сел в штангу, и тут он пригласил его в кастрюлю и сказал: "Двадцать месяцев не хватает".