如何利用多 GPU、多服务器分布式训练深度学习模型,本篇介绍 pytorch 分布式训练相关。

一些概念

  • node 物理节点,即服务器、机器。注意每一个节点可以有多个 GPU,即一机多卡
  • rank 进程序号,常用于进程通信。每个进程对应一个 rank
  • local_rank 相对于 rank (整个分布式集群而言),local_rank 表示一台机器内部的相对进程序号。rank 和 locak_rank 独立,当是单机多卡时,rank 等同于 local_rank
  • nnodes 物理节点个数
  • node_rank 物理节点序号
  • nproc_per_node 每个物理节点上面运行的进程数,常对应想要运行的 GPU 个数,因此取值小于物理节点上 GPU 个数。
  • group 进程组,默认只有一个组
  • world size 全局并行数,即一个分布式任务中,全局 rank 的数量。

例子:
每个 node 包含 8 个 GPU,且 nproc_per_node=4, nnodes=5,机器的 node_rank=10,请问 world_size 是多少?
全局并行数 world size 等于真实的进程数,即 nproc_per_node x nnodes = 4 x 5 = 20,其他没有使用的 GPU 是没有参与计算,不算并行数。

pytorch 分布式特点

torch.distributed 软件包和 torch.nn.parallel.DistributedDataParallel 模块由全新的、重新设计的分布式库提供支持。新的库的主要亮点有:

  • 新的 torch.distributed 是性能驱动的,并且对所有后端 (Gloo,NCCL 和 MPI) 完全异步操作
  • 显着的分布式数据并行性能改进,尤其适用于网络较慢的主机,如基于以太网的主机
  • 为 torch.distributedpackage 中的所有分布式集合操作添加异步支持
  • 在 Gloo 后端添加以下 CPU 操作:send、recv、reduce、all_gather、gather、scatter
  • 在 NCCL 后端添加barrier操作
  • 为 NCCL 后端添加new_group支持

torch.distributed 包提供了一个启动实用程序 torch.distributed.launch,此帮助程序可用于为每个节点启动多个进程以进行分布式训练,它在每个训练节点上产生多个分布式训练进程。

单节点多 GPU

1
2
# 这里 NUM_GPUS_YOU_HAVE 为你服务器上 GPU 的个数,建议使用同型号 GPU 进行并行计算
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

多节点多 GPU

以两个节点为例,更多节点可以类推。假设节点1(IP: 192.168.1.100)为主节点,节点2(IP: 192.168.1.200)为从节点。
在节点1上

1
2
3
# 因为是2个节点,所以 nnodes = 2
# 因为节点1是主节点,所以指定节点1 的IP
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=0 --master_addr="192.168.1.100" --master_port=8888 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

在节点2上

1
2
# 在每个从节点上,指定主节点IP地址和通信端口号
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=1 --master_addr="192.168.1.100" --master_port=8888 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

与单节点单GPU的区别:

  • 后端最好用 NCCL,才能获取最好的分布式性能
  • 训练代码必须从命令行解析 --local_rank=LOCAL_PROCESS_RANK
    1
    2
    3
    4
    5
    6
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int)
    args = parser.parse_args()

    torch.cuda.set_device(arg.local_rank)
  • 分布式初始化
    1
    torch.distributed.init_process_group(backend='nccl', init_method='env://')
  • 模型分布式
    1
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[arg.local_rank], output_device=arg.local_rank)

并行

深度模型训练加速建议:

  1. 当模型不大时(可以加载到一台服务器上时),使用数据并行;
  2. 当模型较大时,可将模型分解放在多个服务器上运行,即模型并行;

参考文献

  1. PyTorch多卡/多GPU/分布式DPP的基本概念(node&rank&local_rank&nnodes&node_rank&nproc_per_node&world_size)
  2. pytorch 1.0 分布式