Российские ученые в десятки раз сократили расход памяти графовых нейросетей

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

Исследователи из Yandex Research решили эту проблему на самом низком уровне программной оптимизации — непосредственно в операциях, выполняемых графическими процессорами. Они проанализировали наиболее ресурсоемкие вычислительные этапы графовых нейросетей и разработали новые реализации базовых GPU-операций, которые эффективнее используют вычислительные блоки и локальную память видеоускорителей, значительно сокращая количество обращений к более медленной внешней памяти.

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

Испытания показали, что новые программные компоненты ускоряют отдельные графовые модели до 8,5 раза, а некоторые операции — в 3,9–10 раз. Одновременно пиковое потребление памяти сокращается до 76 раз, что позволяет размещать в памяти графических процессоров существенно более крупные графы и масштабировать обучение моделей.

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