Библиотека YaFSDP (Yet Another Fully Sharded Data Parallel) помогает оптимизировать использование большими языковыми моделями ресурсов большого числа графических процессоров на всех этапах обучения: предварительном, с учителем, выравнивании.

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

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

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

Решение было представлено «Яндексом» в июне 2024 года. Доступно как открытое ПО.

Репозиторий на GitHub

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

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

• Смотреть