Учёные Яндекса ускорили графовые нейросети в 8,5 раза и снизили память в 76 раз

Графовые нейросети широко применяются в рекомендательных системах, поиске, социальных сетях, транспорте и финансах, но их производительность часто ограничена скоростью обмена данными между GPU и памятью. Российские учёные проанализировали самые ресурсоёмкие операции и создали их эффективные реализации, реже обращающиеся к внешней памяти. В экспериментах отдельные модели ускорились в 3,9–10 раз, а пиковое потребление памяти снизилось до 76 раз.

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

«Наше решение эффективно работает на системах любого размера, так как затрагивает самый близкий к железу уровень оптимизации — ускорение операций, исполняемых на GPU. Поэтому ускорение происходит в рамках одной GPU, и на системе из тысяч — ускорение будет на каждой из них. К тому же, наше решение позволяет обучать модели, используя меньше GPU по сравнению со стандартными методами», - пояснил исследователь Yandex Research Фёдор Великонивцев.