Подробнее о тренировке OPT — огромной языковой модели с 175B параметров от Meta AI

Чем глубже погружаешься в методы тренировки больших моделей, тем больше понимаешь насколько там боли. С одной стороны у тебя может быть 1000 GPU, но как заставить их работать вместе и что делать когда модель начинает расходиться?

Для обучения OPT использовали коктейль из
1. MegatronDeepSpeed — эффективный tensor parallel
1. NVIDIA Apex — mixed precision
1. Fairscale — fully-sharded data parallel и элистичность тренировки, т.е. чтобы когда умирала какая-то GPU вся тренировка не останавливалась

Модель максимально близка к GPT3. Та же максимальная длина в 2048 токенов, тот же токенизатор от GPT2, ReLU. Pre-norm или post-norm в статье на написано, но предполагаю что pre-norm. Использовали обычный fp16 для весов модели, а ADAM был в fp32.

Данные это смесь данных RoBERTa (BookCorpus, Stories, CCNews), часть The Pile, и PushShift.io Reddit. Дедуплицировали документы с помощью Min-hashLSH, сказали что в The Pile очень много дубликатов. В конце-концов получилось 180B токенов.

Самое весёлое: процесс тренировки. Тренировали самую большую модель 2 месяца, её приходилось вручную перезапускать 35 раз. Автоматически она перезапустилась ещё 70. То есть в среднем чаще чем раз в день с тренировкой что-то происходило. При ручном перезапуске делали диагностику hardware и выключали сломанные ноды. Когда лосс начинал резко расти, модель откатывали до последнего чекпоинта и уменьшали LR. Кроме этого в начале тренировки использовали очень маленький gradient clipping в 0.3.

Вот когда я разчувствовался при чтении статьи это на секции где авторы описывают что во время тренировки они пытались менять оптимизатор в ADAMW на SGD (не помогло, вернули обратно), ресетить скейлинг лосса (это когда мы лосс умножаем на некоторое число, чтобы подсчёт градиентов был более численно стабильным, важно для fp16) и даже менять версию Megatron. У кого-то были напряжённые недели.

Кажется 2022 запомнится демократизацией доступа к большим моделям, что может привести к куче ноывых интересных статей как эти модели применять на практике и какие у них есть косяки. Stay tuned.