Недавно Яндекс выпустил в опенсорс свой новый инструмент — YaFSDP, который поможет значительно ускорить обучение LLM (больших языковых моделей).
Представьте, что вы тренируете модель на одной GPU: прогоняете данные через сеть, считаете потери, делаете обратное распространение ошибки и обновляете веса модели…
А теперь представьте, что у вас 4 GPU. И тут начинаются проблемы: нужно синхронизировать работу всех GPU и пересылать огромное количество данных между ними. Это занимает кучу времени и съедает всю память. Например, для модели Llama 70B в fp16 нужно переслать 280 GB данных на каждую итерацию!
Поэтому на помощь приходят методы вроде ZeRO и FSDP, которые помогают эффективно разрезать и распределять веса, градиенты и состояния оптимизатора между всеми GPU. Но и у них есть свои недостатки: ZeRO может замедлять обучение на больших кластерах, а FSDP иногда выделяет больше памяти, чем нужно, и тормозит из-за коммуникаций между GPU.
Яндекс взял лучшее от этих методов и добавил свои улучшения. YaFSDP экономит память, избегает ненужных коммуникаций и оптимизирует процессы так, чтобы всё работало гладко и быстро. Они создали буферы для хранения данных, используют cuda-стримы для параллельных вычислений и коммуникаций, и избавились от проблем с «ёлочкой» (долгие задержки из-за пересылки данных).
Освобождается память, чтобы не делать лишние чекпоинты активаций и ускорить обучение на 25%!
Подход позволяет быстро пересылать данные между GPU, не создавая лишних задержек.
YaFSDP интегрируется с torch и не требует сложной настройки.
Но и здесь есть ограничения. Наибольшая производительность достигается, если слои вызывались так, чтобы соответствующие им буферы чередовались. Также метод подразумевает, что с точки зрения оптимизатора может быть только одна группа весов с большим числом параметров.
Исходный код YaFSDP здесь.