Эффективное масштабирование графовых нейронных сетей за счёт реализации алгоритмов, учитывающих устройство графических ускорителей

Эффективное масштабирование графовых нейронных сетей за счёт реализации алгоритмов, учитывающих устройство графических ускорителей

Великонивцев Ф.С. (НИУ ВШЭ, Москва, Россия)

Аннотация

Производительность графовых нейронных сетей (GNN) ограничена разреженным нерегулярным доступом к памяти. Популярные библиотеки, такие как DGL и PyTorch Geometric, поддерживают обобщённую передачу сообщений, однако сложные слои нередко материализуют промежуточные тензоры по рёбрам, увеличивая трафик памяти и ограничивая масштабируемость на больших графах. Мы применяем подход, ориентированный на операции ввода-вывода и арифметическую интенсивность, и показываем, что широко используемые слои классифицируются на три семейства: свёртки на основе умножения разреженной матрицы на плотную (Sparse-Dense Matrix Multiplication, SpMM), агрегации на основе редукции и слои на основе внимания (GATv2/граф-трансформер). Для каждого семейства мы разрабатываем ядра для GPU, снижающие перемещение данных, улучшающие локальность и остающиеся устойчивыми на реалистичных графах. Мы также исследуем перестановку вершин графа и обнаруживаем, что её влияние зависит от отображения ядра: она устойчивее помогает ядрам с параллелизмом по соседям (в которых доминирует gather), чем проектам с параллелизмом по признакам. Эмпирически наши слитые ядра внимания достигают ускорения до 3,9 раз для граф-трансформера (медиана 1,6 раз), а варианты с тензорными ядрами Tensor Cores (блочно-разреженные) – до 7,3x на локально плотных графах; для слоя GATv2 мы достигаем ускорения до 8,5 раз (медиана 2,0 раз) при снижении пикового потребления памяти до 76 раз (медиана 6 раз). Наши ядра редукции с учётом степеней достигают ускорения до 10 раз (медиана 2,6 раз). Для слоёв на основе SpMM корректно кэшированный cuSPARSE (библиотека NVIDIA для разреженной линейной алгебры) достигает ускорения до 8 раз над DGL (самой популярной библиотекой для графового машинного обучения Deep Graph Library) и превосходит оцениваемые пользовательские базовые решения в большинстве оценок. Мы выпускаем наши реализации как готовые замены для поддержки воспроизводимого, учитывающего оборудование ускорения GNN.

Ключевые слова

графовые нейронные сети; ядра, учитывающие ввод-вывод; ускорение на GPU; умножение разреженной матрицы на плотную; слои на основе внимания; слияние ядер; ограниченные памятью вычисления; арифметическая интенсивность; перестановка вершин графа.

Издание

Труды Института системного программирования РАН, том 38, вып. 5, 2026, стр. 195-216.

ISSN 2220-6426 (Online), ISSN 2079-8156 (Print).

DOI: 10.15514/ISPRAS-2026-38(5)-12

Для цитирования

Великонивцев Ф.С. Эффективное масштабирование графовых нейронных сетей за счёт реализации алгоритмов, учитывающих устройство графических ускорителей. Труды Института системного программирования РАН, том 38, вып. 5, 2026, стр. 195-216. DOI: 10.15514/ISPRAS-2026-38(5)-12.

Полный текст статьи в формате pdf Вернуться к содержанию тома