Яндекс выпустил YaFSDP: инструмент для ускорения обучения больших языковых моделей

Яндекс разработал и выложил в открытый доступ инструмент YaFSDP, который призван ускорить обучение больших языковых моделей (LLM) и сократить расходы на графические процессоры (GPU).

Яндекс выпустил YaFSDP

Проблемы при обучении на нескольких GPU

Обучение на нескольких GPU может значительно ускорить процесс, но не все так просто!

Какие сложности могут возникнуть?

• ⏱ Синхронизация: Обмен данными между GPU может тормозить обучение.
• 🗃️ Памяти мало: Каждая GPU имеет ограниченную память, а данные нужно делить на части.
• ⚙️ Настройка: Управление кластером GPU и согласование параметров обучения могут быть сложными.
• 💰 Дорого: Много GPU - это дорогое удовольствие, и обучение потребляет много энергии.

YaFSDP

Почему важно оптимизировать память при обучении языковых моделей ?Давайте разберемся!

Память потребляют:

• Веса, градиенты, оптимизатор: Их объём зависит от числа процессов и стремится к нулю при увеличении их количества.
• Буферы: Занимают постоянный объём памяти.
• Активации: Самый "прожорливый" компонент! Зависит от размера модели и количества токенов.

Выходит, единственное, что занимает память, это активации.

Из этого следует такое решение: Чекпоинт активаций!

1) Сохраняем только активации между блоками трансформера.

2) Перевычисляем их на backward, экономя память (до 5 GB).

3) Но теряем 25% времени на обучение.

Подробнее про YaFSDP рассказал Михаил Хрущев NLP developer in Yandex в этой статье : https://habr.com/ru/companies/yandex/articles/817509/