Авторы

Бартломей Гаврих, Кацпер Пьеткун, Хешам Мостафа

Введение

В области графовых нейронных сетей (GNN) обучающие модели на больших графах создают проблемы с точки зрения ограничений памяти и эффективной связи в распределенных средах. Для решения этих проблем была представлена ​​SAR, чистая библиотека Python, созданная поверх DGL (библиотека глубокого графа). SAR обеспечивает распределенное обучение GNN на больших графах, поддерживая как полное пакетное обучение, так и обучение на основе выборки. В полнопакетном режиме SAR обучает модель на всем графе, используя все узлы и ребра. В режиме выборки SAR обучает модель на подмножестве узлов и ребер, выбранных из графа. Оба режима имеют свои преимущества и недостатки, в зависимости от размера и плотности графа, архитектуры модели и задачи обучения. В этом сообщении блога мы рассмотрим ключевые функции SAR, включая построение графиков, загрузку данных и полнофункциональное обучение.

SAR реализует новую схему под названием Sequential Aggregation and Rematerialization (SAR) для комплексного обучения. SAR снижает пиковое потребление памяти на машину и обеспечивает линейное уменьшение потребления памяти моделью на одного работника с количеством рабочих. Это достигается за счет устранения большей части избыточности данных, связанной со стандартным распределенным обучением из-за эффекта ореола. Эффект ореола относится к тому факту, что каждому воркеру необходимо хранить не только свой собственный раздел графа, но и некоторые перекрывающиеся узлы и ребра из других разделов для выполнения передачи сообщений. SAR позволяет избежать этого, последовательно агрегируя сообщения из разных разделов графа и повторно материализуя их по запросу.

Построение графика

SAR предоставляет функции для загрузки данных и построения графиков для облегчения эффективного обучения GNN. После разделения графа с помощью функции DGL partition_graph() SAR может загрузить данные графа с помощью функции sar.load_dgl_partition_data(), которая возвращает объект PartitionData. Этот объект можно использовать для создания различных типов графоподобных объектов для передачи в модели GNN. Объект PartitionData. Этот объект можно использовать для построения различных типов графоподобных объектов для передачи в модели GNN.

Для полного пакетного обучения SAR позволяет построить объект распределенного графа с именем GraphShardManager с помощью функции sar.construct_full_graph(). GraphShardManager инкапсулирует N объектов графа DGL, называемых GraphShards в каждом рабочем потоке, где N — количество разделов/воркеров. iᵗʰ GraphShard в GraphShardManager at worker j содержит ребра, соединяющие узлы в разделе iк узлам в разделе j. GraphShardManager реализует такие методы, как update_all() и apply_edges(), обеспечивающие последовательное агрегирование и повторную материализацию для распределенных прямых и обратных проходов во время обучения.

Пример

Чтобы наглядно представить, как работает SAR, мы будем использовать небольшой пример графа, который показан в левой части рисунка 1. Сначала мы представим процесс разбиения графа DGL (разделение большого графа на более мелкие подграфы для распределенного обучения), а затем Рабочий процесс САР. Чтобы упростить задачу, подготовим граф к обучению на трех узлах, разделив его на три раздела. Алгоритм разделения DGL равномерно распределяет узлы и ребра графа по разделам, обеспечивая сбалансированную нагрузку на каждого работника во время обучения.

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

После разделения графа каждый рабочий создает GraphShardManager, который состоит из N GraphShards. GraphShard представляет собой двудольный граф, содержащий ребра, соединяющие узлы в двух разделах (эти два раздела могут быть одинаковыми). Двухкомпонентная структура GraphShard обеспечивает эффективную передачу и агрегирование сообщений между узлами в двух разделах. На рис. 3 представлен процесс создания GraphShards для первого раздела секционированного примера графа. Обратите внимание, что существует три объекта GraphShard, по одному на каждый раздел.

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

После создания GraphShards для каждого раздела они сохраняются и управляются GraphShardManager. GraphShardManager предоставляет интерфейс для доступа к сегментам GraphShard и управления ими, например методы для передачи сообщений, агрегирования и обновления функций узлов и краев. GraphShardManager играет решающую роль в обеспечении эффективного распределенного обучения на разделенном графе за счет синхронизации операций между рабочими процессами. На рис. 5 представлен окончательный вид трех GraphShardManager, созданных для примера графа.

Обучение с САР

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

Чтобы облегчить обучение SAR с полным графом, в существующей модели GNN могут потребоваться некоторые модификации. Edge softmax, критический компонент в слоях GNN, основанных на внимании, несовместим с графическими объектами SAR. Однако SAR предоставляет sar.edge_softmax() в качестве альтернативы для решения этой проблемы. Кроме того, последовательная рематериализация SAR во время обратного прохода требует явного знания любых обучаемых параметров, используемых в функции сообщения. Декоратор sar.core.message_has_parameters() информирует SAR об этих параметрах и включает точное обратное распространение градиента.

Анимация, показанная на рисунке 6, представляет процесс передачи сообщений с использованием SAR для примера графика. Его можно разбить на несколько пунктов:

  1. Локальная агрегация — агрегация функций, которые хранятся внутри воркера (нет необходимости во внешней связи)
  2. Удаленная выборка — связь с другими работниками для получения необходимых функций узлов.
  3. Удаленная агрегация — агрегация полученных данных

Обратите внимание, что в режиме по умолчанию, без каких-либо изменений, SAR будет извлекать и агрегировать данные от других рабочих процессов последовательно один за другим, что приводит к материализации не более 2 разделов графа на каждой машине в любой момент времени. Таким образом, снижается потребность в памяти.

SAR поддерживает три режима обучения: «Последовательное агрегирование и повторная материализация» (Режим 1), «Последовательное агрегирование» (Режим 2) и «Однократное агрегирование» (Режим 3). Каждый режим предлагает уникальный компромисс между эффективностью использования памяти и коммуникационными издержками. Режим 1, который используется по умолчанию, сводит к минимуму потребление памяти на одного рабочего и допускает линейное сокращение с количеством рабочих. Режим 2 строит вычислительный граф во время прямого прохода, поэтому нет необходимости в последовательной повторной материализации. Режим 3 направлен на минимизацию раундов связи путем агрегирования данных за один раз.

Основной объект распределенного графа в SAR, GraphShardManager, обеспечивает замену собственного объекта DGL DGLGraph, обеспечивая простую интеграцию с существующими слоями GNN. GraphShardManager реализует распределенные версии методов update_all() и apply_edges(), необходимых для передачи сообщений и операций с графами. SAR прост в использовании и требует минимальных модификаций кода обучения с одним узлом для запуска распределенного обучения. На рис. 7 представлены необходимые изменения, которые необходимо включить в код. Это:

  • Инициализация связи между работниками
  • Загрузка раздела графика и извлечение необходимой информации, такой как функции и метки.
  • Создание объекта GraphShardManager (замена собственного графа DGL)
  • Синхронизация параметров модели по всем воркерам перед началом обучения
  • Синхронизация рассчитанных градиентов по всем воркерам после каждой обратной операции

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

Полученные результаты

Оценка SAR включала тестирование его производительности на двух больших графиках из наборов данных Open Graph Benchmarks (OGB): ogbn-products и ogbn-papers100M. Эксперименты проводились на кластере машин с процессорами Xeon, где каждый процессор имел 36 физических ядер, а каждая машина имела 256 ГБ оперативной памяти.

Во время оценки GraphSage на графике ogbn-products SAR продемонстрировал сравнимые накладные расходы на связь с параллельным обучением домена, предлагая при этом небольшое преимущество во времени выполнения и повышение эффективности использования памяти по мере увеличения количества разделов. Для модели GAT SAR был медленнее, чем собственное параллельное обучение из-за накладных расходов на связь во время обратного прохода, но, что примечательно, SAR потреблял гораздо меньше памяти по сравнению со стандартным параллельным обучением с увеличивающимся числом рабочих процессов.

При разбиении графа ogbn-papers100M на 32, 64 и 128 машин SAR продемонстрировал значительные преимущества памяти по сравнению с собственным параллельным обучением (рис. 8). SAR был почти в четыре раза более эффективным с точки зрения использования памяти, чем параллельное обучение домена при обучении GAT на 128 машинах. Напротив, встроенное параллельное обучение сталкивалось с ошибкой нехватки памяти (OOM) при попытке обучить GAT на 32 машинах. Однако дополнительные коммуникационные издержки SAR повлияли на его поведение при масштабировании во время выполнения по мере увеличения количества машин. Увеличение числа рабочих с 64 до 128 увеличило время выполнения SAR в эпоху, поскольку обучение стало более тесно связано с общением. Параллельное обучение предметной области на GAT со 128 рабочими процессами было значительно быстрее из-за меньшего количества коммуникаций.

Для GraphSage, где объем обмена данными оставался одинаковым для всех методов, SAR имел небольшое преимущество во время выполнения. SAR удалось вдвое сократить потребление памяти при обучении сети GraphSage на 128 машинах (рис. 9).

Краткое содержание

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

Ссылки