SmartNIC-ускорение обучения GNN: система SNI-GNN сокращает трафик до 45% и ускоряет вычисления до 3,6 раз

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

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

Помимо предсказания, в SNI-GNN применяются два дополнительных механизма: выборка граничных узлов по важности и асинхронный конвейер DPU-GPU с повторным использованием промежуточных результатов. Такой подход снижает нагрузку на каналы связи и ускоряет общий цикл обучения без ущерба для качества модели.

Эксперименты проводились на оборудовании NVIDIA BlueField-3. Система интегрируется с современными фреймворками полнографового обучения и показывает сокращение коммуникации на 21–45%. По сравнению с базовым решением BNS-GCN она обеспечивает ускорение в 1,3–3,6 раза, а относительно системы SANCUS — до 1,29 раза. При этом потеря точности не превышает 0,01.

Масштабируемость SNI-GNN проверена на конфигурации до 16 GPU и графах с десятками миллионов рёбер. По словам авторов, полученные результаты демонстрируют, что предсказание данных внутри сети на базе SmartNIC является практичным дополнением к методам разбиения графа и сжатия данных для эффективного обучения GNN в больших масштабах.