DGL分布式GNN技术分享
1.背景
图是一种由节点和边组成的数据结构。在现代社会，图结构数据无处不在，如分子数据、社交网络等。图神经网络（GNN）作为分析这类数据的强有力工具，正受到越来越多的关注。随着数据规模的增长，传统的单机单卡GNN已无法满足需求。分布式图神经网络通过在多个设备上并行处理数据，显著提高了数据处理的效率。DGL（Deep Graph Library）采用数据并行策略，实现了分布式GNN的训练。
数据并行化分布式GNN训练流程：
数据划分：将数据分配到多个设备上，可采用随机分配或metis等图划分算法。 模型初始化：使用相同的随机种子，对模型进行初始化，保证模型的初始参数完全一致。 训练阶段： 数据采样：各设备独立采样训练节点。 模型训练：计算损失函数并反向传播得到各个设备上的梯度。 梯度同步：将不同设备产生的梯度信息共享到其他设备，确保不同设备的梯度一致。 梯度更新：使用梯度对模型参数进行更新，由于不同设备上的初始参数一致，每一次更新的梯度也一致，所以模型参数始终保持一致。 测试阶段&评价阶段：使用单台设备进行测试和评价。
2.基于DGL的数据并行化分布式GNN训练
数据划分过程：
如果单个设备能够存储完整图数据，则复制到每台设备并分配训练节点；否则，使用metis等图划分算法进行数据切分。
metis算法介绍：
metis算法是一种图划分算法，主要用于将一个图划分成若干个部分，使得每个部分的节点数尽可能接近，并且跨部分的边数量尽可能少。
metis算法的核心思想是通过迭代的方法，逐步优化图的划分。它采用递归二分的方法。在每次迭代中，算法会选择一个节点作为”分割点“，然后尝试将这个点所在的部分与其他部分分开。分割点的选择基于多种因素，包括节点的度、节点的连接性等。
在下图中，我们展示了一个用metis算法将图划分成4个part的结果。可以发现metis算法尽可能的将有边的节点分到了同一个子图中。
主要组件：
采样器(Sampler)：DGL基于远程流程调用技术（RPC）实现了分布式采样。用户可以在DGL使用与单卡场景下相同的api调用分布式采样器进行邻居采样。采样生成的小批量子图会被送回到训练器中，供训练器使用。 训练器(Trainer)：数据并行化分布式GNN训练过程中会有许多训练器。对于每个训练器，负责使用从采样器中获得的小批量子图和从嵌入存储器中获得的数据计算得到更新的梯度信息。 服务器(Server)：主要负责梯度同步和存储图数据。在DGL使用PyTorch后端的情况下，使用PyTorch的all-reduce进行地图同步。其他后端（MXNet和TensorFLow）下使用服务器同步梯度。 嵌入存储器(Embedding/Value Memory)：服务器的一个子模块。负责管理节点和边的数据，可以通过api取出和写入数据。
下图展示了DGL分布式图神经网络的整体训练流程，图结构数据和嵌入信息，被保存在服务器中。采样器会根据图结构获取mini-batch子图，然后将子图发送给训练器。训练器可以通过访问服务器中的嵌入存储器获取所需要的节点和边的信息。经过训练后，训练器会将产生的梯度信息发送回服务器，服务器会将梯度信息进行同步。
3.技术特点
数据并行化
采用数据并行化的方式，将训练数据直接分配给每个训练器，避免了训练过程中的复杂数据通信。只需要在数据采样和梯度同步的过程中进行通信，极大的简化了通信过程。
可扩展性
采用数据并行化技术，能够充分的利用多块显卡上的更多的计算资源，能够更高效的在大规模图数据上进行训练。
4.在OpenHGNN上的实践
线程划分
产生多个线程，线程的数量和要使用的gpu的数量相等。
import torch.multiprocessing as mp
mp.spawn(self.distributed_run, nprocs=num_gpus)
分布式初始化
使用pytorch的分布式初始化api进行初始化。
num_gpus=len(self.config.gpu)
torch.distributed.init_process_group(
backend="gloo",
init_method=f"tcp://127.0.0.1:12391",
world_size=num_gpus,
rank=proc_id,
)
分配gpu
为每个线程分配gpu。
self.config.gpu = self.config.gpu[proc_id]
模型训练
训练过程的代码基本和单卡的训练代码相同。
if self.config.gpu < 0:
self.config.device = torch.device('cpu')
elif self.config.gpu >= 0:
if not torch.cuda.is_available( ):
self.config.device = torch.device('cpu')
warnings.warn("cuda is unavailable, the program will use cpu instead. please set 'gpu' to -1.")
else:
self.config.device = torch.device('cuda', int(self.config.gpu))
# test add profiler
if self.config.mini_batch_flag == False:
if hasattr(self.config, 'max_epoch'):
self.config.max_epoch = self.config.max_epoch // num_gpus
if hasattr(self.config, 'line_profiler_func'):
from line_profiler import LineProfiler
prof = LineProfiler
for func in self.config.line_profiler_func:
prof = prof(func)
prof.enable_by_count()
self.config.logger = Logger(self.config)
set_random_seed(self.config.seed)
trainerflow = self.specific_trainerflow.get(self.config.model, self.config.task)
if type(trainerflow) is not str:
trainerflow = trainerflow.get(self.config.task)
if self.config.hpo_search_space is not None:
# hyper-parameter search
hpo_experiment(self.config, trainerflow)
else:
flow = build_flow(self.config, trainerflow)
result = flow.train()
if hasattr(self.config, 'line_profiler_func'):
prof.print_stats()
return
总结
分布式GNN训练更充分的使用了多级多卡的算力，能够更快的处理大规模的图数据。数据并行化方法在各种分布式GNN训练方法中，是逻辑最清晰、实现最简洁的一种。
实验室算法库简介
两算法库均获得OpenI启智社区“社区优秀孵化项目奖”
GammaGL：
GammaGL是一个基于TensorLayerX实现的开源的多框架的图神经网络（GNN）算法平台，该框架抽象出图数据、消息传递、采样等组件，已经实现了二十余种图数据集和四十余种术界经典GNN算法。 GammaGL发布0.5版本 https://github.com/BUPT-GAMMA/GammaGL
OpenHGNN：
OpenHGNN是一个基于深度学习框架PyTorch和图神经网络框架DGL作为底层框架的异质图神经网络（HGNN）工具包。该工具包专注于HGNN的算法模型，抽象出了数据集管理、算法流程、任务评测等模块，提供了超参数调优组件和最优的超参数设置，发布了HGNN设计空间和Benchmark等。OpenHGNN发布0.7版本OpenHGNN：https://github.com/BUPT-GAMMA/OpenHGNN
