7.2 启动分布式训练/推理的工具

DGL 在 dgl/tools 目录下提供了一个启动脚本 launch.py,用于在集群中启动分布式训练作业。该脚本做出以下假设:

  • 分区后的数据和训练脚本已提供给集群或所有工作机器都可以访问的共享存储(例如 NFS)。

  • 调用 launch.py 的机器对所有其他机器具有免密码 SSH 访问权限。启动机器必须是其中一台工作机器。

下面展示了在集群中启动分布式训练作业的示例。

python3 tools/launch.py               \
  --workspace /my/workspace/          \
  --num_trainers 2                    \
  --num_samplers 4                    \
  --num_servers 1                     \
  --part_config data/mygraph.json     \
  --ip_config ip_config.txt           \
  "python3 my_train_script.py"

参数指定工作空间路径、分区元数据 JSON 和机器 IP 配置文件的位置,以及在每台机器上启动的训练器(trainer)、采样器(sampler)和服务器(server)进程数量。最后一个参数是启动命令,通常是模型训练/评估脚本。

ip_config.txt 的每一行是集群中一台机器的 IP 地址。可选地,IP 地址后面可以跟着一个网络端口(默认为 30050)。典型的示例如下:

172.31.19.1
172.31.23.205
172.31.29.175
172.31.16.98

启动脚本中指定的工作空间是机器上的工作目录,其中包含训练脚本、IP 配置文件、分区配置文件以及图分区数据。所有文件的路径应指定为相对于工作空间的相对路径。

启动脚本在每台机器上创建指定数量的训练作业(--num_trainers)。此外,用户需要指定每个训练器(trainer)的采样器(sampler)进程数量(--num_samplers)。