前言

工业界的图规模都非常大，少说也是上千万的顶点+上亿的边，单机训练不现实，必须借助多机分布式训练。然而目前主流的图训练框架PyG、DGL对图的多机分布式训练支持都不太好。工业界好像阿里的Euler、百度的PGL可以支持分布式训练。今天介绍一下亚马逊DGL针对分布式训练所做的优化。

摘要

GNN广泛应用在推荐、搜索、风控等领域，在这些领域，图的规模往往非常大，有数以亿计的顶点和万亿的边。为支持大规模图的分布式训练，本文提出了DistDGL，它能以mini-batch的方式在多机上进行分布式训练。DistDGL基于DGL框架，它将图数据分布在多台机器上，并基于数据分布，将计算也分布在多台机器上（owner-compute rule）。DistDGL以同步更新的方式进行训练。为了减小分布式训练的通信开销，DistDGL使用一个高效、轻量的图分割算法对图进行分割，在分割时设计了多个负载均衡约束，使得每个分割的子图达到较好的负载均衡。此外，为了减小跨机器的通信，DistDGL在每个子图中保留了halo nodes（正文会介绍到），并且使用了稀疏embedding更新策略。这些优化策略使得DistDGL在分布式训练时能达到较好的高并行效率和内存可扩展性。实验结果表明，在分布式训练时，随着计算资源的增大，DistDGL的训练速度可以线性增长。在16台机器组成的分布式环境中，DistDGL仅用13秒就可以完成1亿节点+30亿边的一个epoch的训练。DistDGL是DGL的一部分，已开源在：https://github.com/dmlc/dgl/tree/master/python/dgl/distributed。

Continue reading →