博客详情

主页 / 博客详情
blog

使用 DGL v0.9.1 加速十亿级图的划分

图是表示关系数据的普遍方式，许多现实世界的应用（如推荐和欺诈检测）都涉及从海量图中学习。因此，GNN 已成为学习图表示的强大模型家族。然而，在海量图上训练 GNN 具有挑战性，其中一个问题是将图数据分发到集群需要高资源需求。例如，将一个拥有 10 亿个节点和 50 亿条边的随机图划分成 8 个分区，需要一台强大的 AWS EC2 x1e.32xlarge 实例（128 个 vCPU，3.9TB 内存）运行 10 小时才能完成任务。

在最新的 DGL v0.9.1 版本中，我们发布了一个新的流水线，用于预处理、划分和分发数十亿节点或边的图，以便进行分布式 GNN 训练。其核心是一种名为 Chunked Graph Data Format (CGDF) 的新数据格式，它按块存储图数据。新的流水线并行处理数据块，不仅降低了每台机器的内存需求，而且显著加速了整个过程。对于拥有 10 亿节点/50 亿边的相同随机图，使用由 8 台 AWS EC2 x1e.4xlarge 实例（每台 16 个 vCPU，488GB 内存）组成的集群，新的流水线可以将运行时间缩短至 2.7 小时，并将成本降低 3.7 倍。

在这篇博客中，我们将逐步说明如何使用这项新功能划分和分发拥有数十亿节点和边的图。

分布式 GNN 训练基础知识

图数据集通常由图结构以及与节点/边相关的特征组成。如果图是异构的（即，有多种类型的节点或边），不同类型的节点/边可能拥有不同的特征集。在多机集群上训练 GNN 模型首先需要用户划分其输入图，这包括两个步骤：

  1. 运行图划分算法（例如，随机、METIS）将每个节点分配到一个分区。
  2. 将图结构和节点/边特征打乱并分发到拥有该分区的目标机器。

一旦图被划分并准备好，用户就可以使用 DGL 的启动工具启动分布式训练程序，该工具将：

  1. 在每台机器上启动一个主图服务器，将本地图分区加载到内存中。图服务器提供远程过程调用 (RPC) 来执行图采样等计算。用户可以选择启动更多共享主图服务器内存数据的备份图服务器，以提高服务吞吐量。
  2. 在每台机器上启动一个键值存储 (KVStore) 服务器，将本地节点/边特征加载到内存中。KVStore 服务提供 RPC 来获取/更新节点/边特征。
  3. 在每台机器上启动一个或多个训练进程。训练进程通过 PyTorch 的 DistributedDataParallel 组件相互连接。在每次训练迭代中，它们向本地或远程图服务器和 KVStore 服务器发出请求，获取一个 mini-batch 的样本，执行梯度下降，并在下一次迭代之前同步其梯度。

下图描绘了系统架构。更多信息请查阅 DGL 用户指南的分布式训练章节。

dist_train

由于图数据的复杂性，图划分通常在单台机器上运行，这要求机器有足够的内存来容纳整个图和特征以及划分算法的运行时内存使用。例如，一个拥有 10 亿个节点和 50 亿条边，每个节点有 50 个特征的随机图，以 DGL 图格式存储时需要 268GB 内存。使用现有的 dgl.distributed.partition_graph API 划分此图需要一台强大的 AWS EC2 x1e.32xlarge 实例（128 个 vCPU，3.9TB 内存），并且运行 10 小时 — 这对于用户进行大规模 GNN 训练是一个显著的瓶颈。

DGL v0.9.1 通过新的分布式图划分流水线解决了这个问题。具体来说，

  • 我们设计了一种 Chunked Graph Data Format (CGDF) 来按块存储大型图数据，以避免将整个图加载到单台机器中。
  • 我们提供了脚本，可以使用多台机器并行划分和分发分块图，从而降低每台机器的内存需求并加速整个过程。

分块图数据格式

分块图数据集组织成一个数据文件夹，包含以下数据文件：

  • 一个 metadata.json 文件，存储图的元信息，例如图名称、节点/边类型、块大小、块文件路径等。
  • 一系列边索引块文件，存储源节点和目标节点 ID。它们通常是纯文本格式。
  • 一系列节点数据块文件。它们通常是以 NumPy 数组二进制格式存储的数组数据。
  • 一系列边数据块文件。它们也采用 NumPy 数组二进制格式。

这里，我们展示了一个随机社交图（节点为用户，边为关注关系）的文件夹结构和数据文件，其中节点包含两个数据：“feat” 和 “label”。有关格式的完整规范以及如何将数据转换为块的技巧，请查阅文档页面。

//data/random_graph_chunked/
  |-- metadata.json            # metadata JSON
  |-- edge_index/              # edge index chunks
    |-- user:follow:user0.txt  # user-follow-user edges chunk 0
    |-- user:follow:user1.txt  # user-follow-user edges chunk 1
    |-- user:follow:user2.txt  # user-follow-user edges chunk 2
    |-- ...
  |-- node_data/           # node data chunks
    |-- user/              # user nodes have two data: "feat" and "label"
      |-- feat0.npy        # feat chunk 0
      |-- feat1.npy        # feat chunk 1
      |-- feat2.npy        # feat chunk 2
      |-- ...
      |-- label0.npy       # label chunk 0
      |-- label1.npy       # label chunk 1
      |-- label2.npy       # label chunk 2
      |-- ...
  |-- edge_data/           # edge data chunks
    |-- user:follow:user/
       |-- ...

运行分布式划分与分发

第一步是准备一个机器集群来划分图。我们建议集群的总内存大小是图数据大小的 2-3 倍，以容纳所需的运行时内存。接下来是设置共享工作区和软件环境。

  1. 设置一个集群中每个实例都可以访问的共享文件夹（例如，使用 NFS）。确保所有实例都可以相互通过 SSH 访问。这里，我们假设文件夹挂载到 /workspace。
  2. 从 DGL 0.9.x 分支克隆和下载脚本到 /workspace。
    git clone https://github.com/dmlc/dgl.git -b 0.9.x /workspace/dgl
  3. 将分块图数据复制/移动到 /workspace/。这里，我们假设数据文件夹位于 /workspace/random_graph_chunked/。
  4. 创建一个 /workspace/ip_config.txt 文件，其中包含每个实例的 IP 地址。
    # example IP config file of a 4 machine cluster
172.31.19.1
172.31.23.205
172.31.29.175
172.31.16.98

然后我们可以运行一个划分算法将每个节点分配到一个分区。这里，我们选择随机划分算法。

python /workspace/dgl/tools/partition_algo/random_partition.py \
    --in_dir=/workspace/random_graph_chunked/
    --out_dir=/workspace/partition_assign/
    --num_partitions=4

上述脚本只是简单地计算一个节点属于哪个分区。然后，我们将分块图和分区分配都传递给 dispatch_data.py 脚本，以物理地将图数据分割成多个部分并分发到整个集群。

python /workspace/dgl/tools/dispatch_data.py \
    --in-dir=/workspace/random_graph_chunked/
    --partitions-dir=/workspace/partition_assign/
    --out-dir=/workspace/random_graph_dist/
    --ip-config=/workspace/ip_config.txt

最终结果将如下所示。然后我们可以按照此处的说明启动分布式训练。

/workspace/random_graph_dist/
  |-- medatdata.json      # metadata JSON file
  |-- part0/              # partition 0
    |-- graph.dgl         # graph structure of partition 0 in DGL binary format
    |-- node_feat.dgl     # node feature of partition 0 in DGL binary format
    |-- edge_feat.dgl     # edge feature of partition 0 in DGL binary format
  |-- part1/              # partition 1
    |-- graph.dgl         # graph structure of partition 1 in DGL binary format
    |-- node_feat.dgl     # node feature of partition 1 in DGL binary format
    |-- edge_feat.dgl     # edge feature of partition 1 in DGL binary format
  |-- part2/
  ...

注意，这些脚本利用多台机器协同划分和处理数据。因此，新的流水线显著更快。对于拥有 10 亿节点/50 亿边的相同随机图，新的流水线使用由 8 台 AWS EC2 x1e.4xlarge 实例（16 个 vCPU，488GB 内存）组成的集群，可以在 2.7 小时内完成划分（快 3.7 倍）。

延伸阅读

9 月 19 日

博客 新闻 发布

作者：DGLTeam，发表于发布