Выложенная в открытый доступ библиотека YaFSDP (Yet another Fully Sharded Data Parallel) помогает оптимизировать использование большими языковыми моделями ресурсов большого числа графических процессоров на всех этапах обучения: предварительном, с учителем, выравнивании. Решение было представлено «Яндексом» в июне 2024 года.

Оно основано на подходе FSDP (Fully Sharded Data Parallelism), встроенном в широко используемую в машинном обучении библиотеку Torch. В нем множество параметров слоя нейросети можно объединить в один, который и будет разбиваться для обработки графическими процессорами. Однако в YaFSDP для хранения промежуточных весов и градиентов выделены два отдельных буфера.

Кроме того, в библиотеке усовершенствовано вычисление градиентов весов.

YaFSDP сравнили с оригинальным подходом FSDP в ускорении обучения моделей Llama 2 и Llama 3. При использовании небольших батчей датасетов российская библиотека позволяла делать это на 20–26% эффективнее. По заявлениям разработчиков, вкупе с другими оптимизациями памяти YaFSDP повышало скорость обучения на 45%.

В «Яндексе» отмечают, что несмотря на то, что YaFSDP рассчитана в первую очередь на большие языковые модели, использовать ее можно и для других нейросетей — например, предназначенных для создания изображений.

Репозиторий на GitHub
Разбор работы библиотеки на Habr

👤Михаил Хрущев