Обучение и генерация текста в nanoGPT

О чём урок

Обучение нейросети и работа с готовой моделью, это две стороны одной медали: без грамотного цикла обучения модель не сойдётся к разумным потерям, а без продуманной генерации даже хорошо обученная модель будет выдавать скучный или сломанный текст. В этом уроке мы разберём train.py и sample.py построчно: как устроен оптимизатор, откуда берётся расписание скорости обучения, что происходит внутри одной итерации с учётом накопления градиента и смешанной точности, и как метод GPT.generate превращает логиты в последовательность токенов. После урока вы сможете самостоятельно запустить обучение на CPU, разобраться в логе и управлять характером генерации через temperature и top_k.

Неофициальный курс AI University по открытому коду (MIT). В уроке приводится код из karpathy/nanoGPT © Andrej Karpathy, лицензия MIT; комментарии переведены на русский, объяснения написаны нашей командой. Курс не связан с автором кода и не одобрен им.

GPT в nanoGPT обучается оптимизатором AdamW, но с одной важной деталью: не все параметры штрафуются одинаково. Весовые матрицы линейных слоёв и таблицы эмбеддингов это двумерные тензоры, и к ним применяется weight decay, то есть на каждом шаге веса слегка "стягиваются" к нулю, что работает как регуляризация и снижает риск переобучения. А вот смещения (bias) и параметры LayerNorm это одномерные векторы, и для них weight decay отключают: штрафовать единственное число сдвига или масштаба не имеет смысла и на практике скорее мешает обучению, чем помогает.

Деление по p.dim() >= 2 это простой и надёжный критерий: он не требует перечислять имена слоёв, а просто смотрит на форму тензора. Таблицы эмбеддингов (wte, wpe) тоже двумерные, поэтому попадают в группу с weight decay наравне с матрицами линейных слоёв. Последняя деталь: fused AdamW. Это реализация шага оптимизатора, где несколько операций объединены в одно CUDA-ядро, что ускоряет обновление весов на GPU.

Скорость обучения (learning rate) в nanoGPT не постоянна: она растёт линейно в начале обучения, а затем плавно убывает по косинусоиде.

Зачем…

План урока

1. Оптимизатор: две группы параметров и fused AdamW
2. Расписание скорости обучения: зачем нужен прогрев
3. Оценка потерь, чекпойнты и один шаг обучения
4. Смешанная точность, autocast и torch.compile
5. Генерация: GPT.generate и sample.py
6. Как читать лог обучения на shakespeare_char
7. Наш пример: генерация текста на shakespeare_char
8. Попробуйте сами