如何利用多 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

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

torch.cuda.set_device(arg.local_rank)

分布式初始化