Как мы работали над YaFSDP

Недавно запустилась библиотека YaFDP, которая ускоряет обучение больших языковых моделей. Подробнее о ней можно почитать в статье на Хабре. А сегодня руководитель группы претрейна YandexGPT Михаил Хрущев ответил на несколько наших вопросов о том, как создавалась библиотека.

Когда и как началась работа над YaFSDP?

За 2 недели до Нового года в 2023-м. К концу года стало ясно, что некоторые наши оптимизации не дают ускорения, которого мы ожидали. Я пошёл искать причины такого и нашел праздничную «ёлочку».

Ориентировались ли на существующие решения? Были ли какие-то референсы, может?

Сам алгоритм предложили в статье ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. Его реализация была также представлена в фреймворке DeepSpeed. К сожалению, она не очень масштабировалась и содержала много багов. Другая реализация, FSDP, сильно лучше, но и там есть проблемы с аллокацией памяти и коммуникациями. Поэтому я начал разрабатывать свою реализацию, к которой впоследствии подключились и другие ребята из моей команды.

Какие трудности возникали во время работы?

Было несколько технических трудностей. Вообще, от начала реализации до выкатки в обучение прошло чуть более месяца — разработка была более или менее гладкой. Самым сложным было продумать идею реализации, на это потребовалось больше месяца.

Здесь труднее всего было реализовать кода, который должен исполняться сразу после backward (обратного распространения ошибки) по слою. Я неожиданно столкнулся с тем, что torch (фреймворк глубокого обучения) совсем не приспособлен для такой хитрой логики. Но решение тем не менее нашлось.

На этом на сегодня все. Рассказывайте о своём опыте работы с LLM в комментариях!

ML Underhood