На перекрёстках тысячи дорог

Графовые нейросети хорошо работают с данными, устроенными как сеть связей: пользователи и подписки, товары и покупки, атомы и химические связи, дороги. Но на GPU они часто оказываются медленнее и прожорливее, чем должны быть.

Студенты ШАД вместе с исследователями Yandex Research нашли, где именно теряются скорость и память, и собрали набор операций, способных ускорить уже существующие модели без изменения их архитектуры.

Чтобы обновить информацию о каждой вершине, графовая нейросеть собирает данные от её соседей. На бумаге всё просто. Но, например, к одному перекрёстку ведут три дороги, к другому — двадцать; у одного пользователя десяток связей, у другого — миллионы. Нужные данные разбросаны по памяти, поэтому ускорителю приходится постоянно искать их, загружать, складывать во временные массивы и записывать обратно.

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

Авторы разобрали основные операции графовых нейросетей и для каждой нашли свой источник замедления.

Первый случай — графовые свёртки, где агрегация признаков соседей представляет собой простую операцию (например каждая вершина суммирует или усредняет данные соседей), которую можно представить в виде умножения разреженной матрицы на матрицу признаков вершин. Разреженная матрица кодирует, от каких вершин идет сигнал: большинство ячеек матрицы пусты, потому что каждый объект связан лишь с небольшой частью графа.

Оказалось, что быстрее многих специальных решений работает cuSPARSE — библиотека NVIDIA для обработки таких матриц на GPU. Это готовый набор инструментов, уже подогнанных под устройство ускорителя. Команда добавила кеширование служебных данных и автоматический перебор наиболее подходящего алгоритма матричного умножения: библиотеке больше не нужно каждый раз заново строить внутреннее представление графа перед прямым и обратным проходами, а также используется самый оптимальный алгоритм умножения.

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

Самая заметная оптимизация коснулась графового attention. Этот механизм позволяет модели определить, какие соседи важнее: например, сильнее учитывать загруженную магистраль и почти игнорировать пустую боковую улицу.

Обычно модель сначала рассчитывает и сохраняет коэффициент важности для каждого ребра. На большом графе получается гигантская ведомость для всех дорог, подписок или химических связей, которая быстро забивает память GPU.

Авторы перенесли на графы принцип FlashAttention. Новый оператор обрабатывает связи частями и сразу добавляет их вклад в результат, не сохраняя всю ведомость целиком. Это как считать поток прямо на перекрёстке, а не отправлять сведения о каждой машине на отдельный склад.

«Главная проблема эффективности GNN на GPU — не количество операций, а память. Мы убрали лишние промежуточные данные, сбалансировали работу между вершинами и задействовали тензорные ядра там, где это позволяет структура графа. По сути, это перенос IO-aware подхода, уже ставшего стандартом для трансформеров, в Graph ML», — объясняет исследователь Yandex Research Фёдор Великонивцев.

Для attention-слоёв оптимизация дала до 8,5 раза меньшую задержку и до 76 раз сократила пиковое потребление памяти. Все решения собрали в набор GPU-операторов для PyTorch, которые можно поставить вместо стандартных компонентов, не перестраивая модель.

Работу приняли как Spotlight на ICML. Такие вот дорожные работы от мира искусственного интеллекта.

@anti_agi