分布式训练 - 多机多卡 (DDP)

起初为调用大规模的模型训练，单卡GPU是不够使用的，需要借用服务器的多GPU使用。就会涉及到单机多卡，多机多卡的使用。在这里记录一下使用的方式和踩过的一些坑。文中若有不足，请多多指正。

由于分布式的内容较多，笔者准备分几篇来讲一次下深度学习的分布式训练，深度学习的框架使用的是Pytorch框架。

----1.分布式训练的理论基础

----2.GPU训练

----3.单机多卡的使用 DP,DDP

----4.多机多卡的使用 DDP

在前边的文章中已经提到了怎样进行单机单卡和单机多卡进行分布式训练，那可能有小伙伴会有疑问能不能进行多机多卡的训练，答案是，必须是可以的。实际上，在现在很多顶级实验室发布的很多预训练模型都是在多机多卡上训练出来的。需要大量的显存才能处理的了大规模的参数和模型。那么今天这篇文章就来详细说说多机多卡的使用来进行分布式训练。

整个分布式文章系列的链接都在上边啦，有需要的小伙伴点击链接就可以看到其他的知识啦！

在单机多gpu可以满足的情况下, 绝对不建议使用多机多gpu进行训练, 我经过测试, 发现多台机器之间传输数据的时间非常慢, 主要是因为我测试的机器可能只是千兆网卡, 再加上别的一些损耗, 网络的传输速度跟不上, 导致训练速度实际很慢. 我看一个github上面的人说在单机8显卡可以满足的情况下, 最好不要进行多机多卡训练。

在详细介绍多机多卡的使用之前，先说一下基本的使用流程：

Pytorch 中分布式的基本使用流程如下：

1. 在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时初始化
   distributed 包。
2. 如果需要进行小组内集体通信，用 new_group 创建子分组
3. 为数据集创建 Sampler
4. 创建分布式并行模型 DDP(model, device_ids=device_ids)
5. 使用启动工具 torch.distributed.launch 在每个主机上执行一次脚本，开始训练
6. 使用 destory_process_group() 销毁进程组

整个的流程跟单机多卡的使用时非常相似的，具体的差别我们会在后边详细讨论。

1.1 初始化–(对于多机多卡分布式很重要，详细说明)

初始化是在程序刚开始的时候运行，在进行多机多卡的训练的时候，要使用torch.distributed.init_process_group() 进行初始化。

函数原型

torch.distributed.init_process_group(backend, 
                                     init_method=None, 
                                     timeout=datetime.timedelta(0, 1800), 
                                     world_size=-1, 
                                     rank=-1, 
                                     store=None)

backend: 后端, 实际上是多个机器之间交换数据的协议
init_method: 机器之间交换数据, 需要指定一个主节点, 而这个参数就是指定主节点的
world_size: 介绍都是说是进程, 实际就是机器的个数, 例如两台机器一起训练的话, world_size就设置为2
rank: 区分主节点和从节点的, 主节点为0, 剩余的为了1-(N-1), N为要使用的机器的数量, 也就是world_size

函数作用
该函数需要在每个进程中进行调用，用于初始化该进程。在使用分布式时，该函数必须在 distributed 内所有相关函数之前使用。

参数详解

1.1.1 backend：指定当前进程要使用的通信后端

在pytorch的官方教程中提供了以下的通信后端，使用分布式时，在梯度汇总求平均的过程中，各主机之间需要进行通信。因此，需要指定通信的协议架构等。torch.distributed 对其进行了封装。

torch.distributed 支持 3 种后端，分别为 NCCL，Gloo，MPI。各后端对 CPU / GPU 的支持如下所示：

gool后端：

gloo 后端支持 CPU 和 GPU，其支持集体通信（collective Communication），并对其进行了优化。

由于 GPU 之间可以直接进行数据交换，而无需经过 CPU 和内存，因此，在 GPU 上使用 gloo 后端速度更快。

torch.distributed 对 gloo 提供原生支持，无需进行额外操作。

NCCL 后端：

NCCL 的全称为 Nvidia 聚合通信库（NVIDIA Collective Communications Library），是一个可以实现多个 GPU、多个结点间聚合通信的库，在 PCIe、Nvlink、InfiniBand上可以实现较高的通信速度。

NCCL 高度优化和兼容了 MPI，并且可以感知 GPU 的拓扑，促进多 GPU 多节点的加速，最大化 GPU 内的带宽利用率，所以深度学习框架的研究员可以利用 NCCL 的这个优势，在多个结点内或者跨界点间可以充分利用所有可利用的 GPU。

NCCL 对 CPU 和 GPU 均有较好支持，且 torch.distributed 对其也提供了原生支持。

对于每台主机均使用多进程的情况，使用 NCCL 可以获得最大化的性能。每个进程内，不许对其使用的 GPUs 具有独占权。若进程之间共享GPUs 资源，则可能导致 deadlocks。

MPI 后端:

MPI 即消息传递接口（Message Passing Interface），是一个来自于高性能计算领域的标准的工具。它支持点对点通信以及集体通信，并且是 torch.distributed 的 API 的灵感来源。使用 MPI 后端的优势在于，在大型计算机集群上，MPI 应用广泛，且高度优化。

但是，torch.distributed 对 MPI 并不提供原生支持。因此，要使用 MPI，必须从源码编译 Pytorch。是否支持GPU，视安装的 MPI 版本而定。

根据官网的介绍, 如果是使用cpu的分布式计算, 建议使用gloo, 因为表中可以看到
gloo对cpu的支持是最好的, 然后如果使用gpu进行分布式计算, 建议使用nccl, 实际测试中我也感觉到, 当使用gpu的时候, nccl的效率是高于gloo的. 根据博客和官网的态度, 好像都不怎么推荐在多gpu的时候使用mpi

基本原则：
用 NCCL 进行分布式 GPU 训练
用 Gloo 进行分布式 CPU 训练

具体详细的分布式通信包底层内容内容可以查看以下链接：

• https://ptorch.com/docs/8/torch-distributed
• https://www.jianshu.com/p/5f6cd6b50140
• https://zhuanlan.zhihu.com/p/76638962

1.1.2 init_method：指定当前进程组初始化方式

分布式任务中，各节点之间需要进行协作，比如说控制数据同步等。因此，需要进行初始化，指定协作方式，同步规则等。

torch.distributed 提供了 3 种初始化方式，分别为 tcp、共享文件 和 环境变量初始化 等。

推荐使用环境变量初始化，就是在你使用函数的时候不需要填写该参数即可，默认使用环境变量初始化。

环境变量初始化(推荐使用)：

默认情况下使用的都是环境变量来进行分布式通信，也就是指定 init_method=“env://”。通过在所有机器上设置如下四个环境变量，所有的进程将会适当的连接到 master，获取其他进程的信息，并最终与它们握手(信号)。

• MASTER_PORT: 必须指定，表示 rank0上机器的一个空闲端口（必须设置）
• MASTER_ADDR: 必须指定，除了 rank0 主机，表示主进程 rank0 机器的地址（必须设置）
• WORLD_SIZE: 可选，总进程数，可以这里指定，在 init 函数中也可以指定
• RANK: 可选，当前进程的 rank，也可以在 init 函数中指定

配合 torch.distribution.launch 使用(详细使用问题请点击该篇文章)。

使用实例：

Node 1（第一台机器）: (IP: 192.168.1.1, and has a free port: 1234)

>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
           --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
           and all other arguments of your training script)

Node 2 （第二台机器）

>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
           --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
           and all other arguments of your training script)

TCP初始化：

看代码：

TCP 方式初始化，需要指定进程 0 的 ip 和 port。这种方式需要手动为每个进程指定进程号。

import torch.distributed as dist

dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456',
                        rank=rank, world_size=world_size)

说明：

不同进程内，均使用主进程的 ip 地址和 port，确保每个进程能够通过一个 master 进行协作。该 ip 一般为主进程所在的主机的 ip，端口号应该未被其他应用占用。

实际使用时，在每个进程内运行代码，并需要为每一个进程手动指定一个 rank，进程可以分布与相同或不同主机上。

多个进程之间，同步进行。若其中一个出现问题，其他的也马上停止。

使用实例：

Node 1：

python mnsit.py --init-method tcp://192.168.54.179:22225 --rank 0 --world-size 2

Node 2：

python mnsit.py --init-method tcp://192.168.54.179:22225 --rank 1 --world-size 2

共享文件系统初始化：

该初始化方式，要求共享的文件对于组内所有进程可见！

看代码：

import torch.distributed as dist

# rank should always be specified
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
                        world_size=4, rank=args.rank)

说明：

其中，以 file:// 为前缀，表示文件系统各式初始化。/mnt/nfs/sharedfile 表示共享的文件，各个进程在共享文件系统中通过该文件进行同步或异步。因此，所有进程必须对该文件具有读写权限。

每一个进程将会打开这个文件，写入自己的信息，并等待直到其他所有进程完成该操作。在此之后，所有的请求信息将会被所有的进程可访问，为了避免 race conditions，文件系统必须支持通过 fcntl 锁定（大多数的 local 系统和 NFS 均支持该特性）。

说明：若指定为同一文件，则每次训练开始之前，该文件必须手动删除，但是文件所在路径必须存在！

与 tcp 初始化方式一样，也需要为每一个进程手动指定 rank。

使用实例

Node 1：

python mnsit.py --init-method file://PathToShareFile/MultiNode --rank 0 --world-size 2

Node 2：

python mnsit.py --init-method file://PathToShareFile/MultiNode --rank 1 --world-size 2

这里相比于 TCP 的方式麻烦一点的是运行完一次必须更换共享的文件名，或者删除之前的共享文件，不然第二次运行会报错。(不建议使用)

TCP 初始化方式和 Env 初始化方式 实例比较：

TCP初始化方式

代码：

import torch.distributed as dist
import torch.utils.data.distributed

# ......
parser = argparse.ArgumentParser(description='PyTorch distributed training on cifar-10')
parser.add_argument('--rank', default=0,
                    help='rank of current process')
parser.add_argument('--word_size', default=2,
                    help="word size")
parser.add_argument('--init_method', default='tcp://127.0.0.1:23456',
                    help="init-method")
args = parser.parse_args()

# ......
dist.init_process_group(backend='nccl', init_method=args.init_method, rank=args.rank, world_size=args.word_size)

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
net = Net()
net = net.cuda()
net = torch.nn.parallel.DistributedDataParallel(net)

执行方式：

# Node 1 : ip 192.168.1.201  port : 12345
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 0 --word_size 3

# Node 2 : 
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 1 --word_size 3

# Node 3 : 
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 2 --word_size 3

说明

1. 在 TCP 方式中，在 init_process_group 中必须手动指定以下参数
2. rank 为当前进程的进程号
3. word_size 为当前 job 的总进程数
4. init_method 内指定 tcp 模式，且所有进程的 ip:port 必须一致，设定为主进程的 ip:port
5. 必须在 rank==0 的进程内保存参数。
6. 若程序内未根据 rank 设定当前进程使用的 GPUs，则默认使用全部 GPU，且以数据并行的方式使用。
7. 每条命令表示一个进程。若已开启的进程未达到 word_size 的数量，则所有进程会一直等待
8. 每台主机上可以开启多个进程。但是，若未为每个进程分配合适的 GPU，则同机不同进程可能会共用 GPU，应该坚决避免这种情况。
9. 使用 gloo 后端进行 GPU 训练时，会报错。
10. 若每个进程负责多块 GPU，可以利用多 GPU 进行模型并行。如下所示：

class ToyMpModel(nn.Module):
    def init(self, dev0, dev1):
        super(ToyMpModel, self).init()
        self.dev0 = dev0
        self.dev1 = dev1
        self.net1 = torch.nn.Linear(10, 10).to(dev0)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to(dev1)

def forward(self, x):
       x = x.to(self.dev0)
       x = self.relu(self.net1(x))
       x = x.to(self.dev1)
       return self.net2(x)

......
dev0 = rank * 2
dev1 = rank * 2 + 1
mp_model = ToyMpModel(dev0, dev1)
ddp_mp_model = DDP(mp_model)
......

Env 初始化方式

代码

import torch.distributed as dist
import torch.utils.data.distributed

# ......
import argparse
parser = argparse.ArgumentParser()
# 注意这个参数，必须要以这种形式指定，即使代码中不使用。因为 launch 工具默认传递该参数
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

# ......
dist.init_process_group(backend='nccl', init_method='env://')

# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)

# ......
# 根据 local_rank，配置当前进程使用的 GPU
net = Net()
device = torch.device('cuda', args.local_rank)
net = net.to(device)
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)

执行方式

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=0 --master_addr="192.168.1.201" --master_port=23456 env_init.py

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=1 --master_addr="192.168.1.201" --master_port=23456 env_init.py

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=2 --master_addr="192.168.1.201" --master_port=23456 env_init.py

说明：

1. 在 Env 方式中，在 init_process_group 中，无需指定任何参数
2. 必须在 rank==0 的进程内保存参数。
3. 该方式下，使用 torch.distributed.launch 在每台主机上，为其创建多进程，其中:
   nproc_per_node参数指定为当前主机创建的进程数。一般设定为当前主机的 GPU 数量
   nnodes 参数指定当前 job 包含多少个节点
   node_rank 指定当前节点的优先级
   master_addr 和 master_port 分别指定 master 节点的ip:port
4. 若没有为每个进程合理分配 GPU，则默认使用当前主机上所有的 GPU。即使一台主机上有多个进程，也会共用 GPU。
5. 使用 torch.distributed.launch 工具时，将会为当前主机创建 nproc_per_node
   个进程，每个进程独立执行训练脚本。同时，它还会为每个进程分配一个 local_rank
   参数，表示当前进程在当前主机上的编号。例如：rank=2, local_rank=0 表示第 3 个节点上的第 1 个进程。
6. 需要合理利用 local_rank 参数，来合理分配本地的 GPU 资源
7. 每条命令表示一个进程。若已开启的进程未达到 word_size 的数量，则所有进程会一直等待。

1.1.3 初始化rank和world_size,timeout,store

rank ： 指定当前进程的优先级

int 值。表示当前进程的编号，即优先级。如果指定 store 参数，则必须指定该参数。

rank=0 的为主进程，即 master 节点。

world_size ：

该 job 中的总进程数。如果指定 store 参数，则需要指定该参数。

timeout ： 指定每个进程的超时时间

可选参数，datetime.timedelta 对象，默认为 30 分钟。该参数仅用于 Gloo 后端。

store

所有 worker 可访问的 key / value，用于交换连接 / 地址信息。与 init_method 互斥。

这里其实没有多难, 你需要确保, 不同机器的rank值不同, 但是主机的rank必须为0, 而且使用init_method的ip一定是rank为0的主机, 其次world_size是你的主机数量, 你不能随便设置这个数值, 你的参与训练的主机数量达不到world_size的设置值时, 代码是不会执行的。

1.1.4 初始化中一些需要注意的地方

首先是代码的统一性, 所有的节点上面的代码, 建议完全一样, 不然有可能会出现一些问题, 其次, 这些初始化的参数强烈建议通过argparse模块(命令行参数的形式)输入, 不建议写死在代码中, 也不建议使用pycharm之类的IDE进行代码的运行, 强烈建议使用命令行直接运行。（argparse模块不懂怎么使用的童鞋请看这篇文章）

其次是运行代码的命令方面的问题, 例如使用下面的命令运行代码distributed.py:

python distributed.py -bk nccl -im tcp://10.10.10.1:12345 -rn 0 -ws 2 

上面的代码是在主节点上运行, 所以设置rank为0, 同时设置了使用两个主机, 在从节点运行的时候, 输入的代码是下面这样:

python distributed.py -bk nccl -im tcp://10.10.10.1:12345 -rn 1 -ws 2 

一定要注意的是, 只能修改rank的值, 其他的值一律不得修改, 否则程序就卡死了初始化到这里也就结束了。

1.2 DistributedSampler

函数原型

torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None)

参数

dataset
进行采样的数据集

num_replicas
分布式训练中，参与训练的进程数

rank
当前进程的 rank 序号（必须位于分布式训练中）

说明

对数据集进行采样，使之划分为几个子集，不同 GPU 读取的数据应该是不一样的。

一般与 DistributedDataParallel 配合使用。此时，每个进程可以传递一个 DistributedSampler 实例作为一个 Dataloader sampler，并加载原始数据集的一个子集作为该进程的输入。

在 Dataparallel 中，数据被直接划分到多个 GPU 上，数据传输会极大的影响效率。相比之下，在 DistributedDataParallel 使用 sampler 可以为每个进程划分一部分数据集，并避免不同进程之间数据重复。

注意：在 DataParallel 中，batch size 设置必须为单卡的 n 倍，但是在 DistributedDataParallel 内，batch size 设置于单卡一样即可。

使用实例

# 分布式训练示例
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel

dataset = your_dataset()
datasampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, batch_size=batch_size_per_gpu, sampler=datasampler)
model = your_model()
model = DistributedDataPrallel(model, device_ids=[local_rank], output_device=local_rank)

1.3 DistributedDataParallel

函数原型

torch.nn.parallel.DistributedDataParallel(module, 
                                          device_ids=None, 
                                          output_device=None, 
                                          dim=0, 
                                          broadcast_buffers=True, 
                                          process_group=None, 
                                          bucket_cap_mb=25, 
                                          find_unused_parameters=False, 
                                          check_reduction=False)

功能

将给定的 module 进行分布式封装， 其将输入在 batch 维度上进行划分，并分配到指定的 devices 上。

module 会被复制到每台机器的每个 GPU 上，每一个模型的副本处理输入的一部分。

在反向传播阶段，每个机器的每个 GPU 上的梯度进行汇总并求平均。与 DataParallel 类似，batch size 应该大于 GPU 总数。

参数解析

• module

要进行分布式并行的 module，一般为完整的 model

• device_ids

int 列表或 torch.device 对象，用于指定要并行的设备。参考 DataParallel。

对于数据并行，即完整模型放置于一个 GPU 上（single-device module）时，需要提供该参数，表示将模型副本拷贝到哪些 GPU 上。

对于模型并行的情况，即一个模型，分散于多个 GPU 上的情况（multi-device module），以及 CPU 模型，该参数比必须为 None，或者为空列表。

与单机并行一样，输入数据及中间数据，必须放置于对应的，正确的 GPU 上。

• output_device

int 或者 torch.device，参考 DataParallel。

对于 single-device 的模型，表示结果输出的位置。

对于 multi-device module 和 GPU 模型，该参数必须为 None 或空列表。

• broadcast_buffers

bool 值，默认为 True

表示在 forward() 函数开始时，对模型的 buffer 进行同步 (broadcast)

• process_group

对分布式数据（主要是梯度）进行 all-reduction 的进程组。

默认为 None，表示使用由 torch.distributed.init_process_group 创建的默认进程组 (process group)。

• bucket_cap_mb

DistributedDataParallel will bucket parameters into multiple buckets so that gradient reduction of each bucket can potentially overlap with backward computation.

bucket_cap_mb controls the bucket size in MegaBytes (MB) (default: 25)

• find_unused_parameters

bool 值。

Traverse the autograd graph of all tensors contained in the return value of the wrapped module’s forward function.

Parameters that don’t receive gradients as part of this graph are preemptively marked as being ready to be reduced.

Note that all forward outputs that are derived from module parameters must participate in calculating loss and later the gradient computation.

If they don’t, this wrapper will hang waiting for autograd to produce gradients for those parameters.

Any outputs derived from module parameters that are otherwise unused can be detached from the autograd graph using torch.Tensor.detach. (default: False)

• check_reduction

when setting to True, it enables DistributedDataParallel to automatically check if the previous iteration’s backward reductions were successfully issued at the beginning of every iteration’s forward function.

You normally don’t need this option enabled unless you are observing weird behaviors such as different ranks are getting different gradients, which should not happen if DistributedDataParallel is correctly used. (default: False)

注意

1. 要使用该 class，需要先对 torch.distributed 进行初进程组始化，可以通过
   torch.distributed.init_process_group() 实现。
2. 该 module 仅在 gloo 和 nccl 后端上可用。
3. 根据分布式原理，Constructor 和 differentiation of the output (或 a function of
   the output of this module) 是一个分布式同步点。在不同的进程执行不同的代码时，需要考虑这一点。
4. 该 module 假设，所有的参数在其创建时，在模型中已经注册，之后没有新的参数加入或者参数移除。对于 buffers也是一样。(这也是由分布式原理决定)
5. 该 module 假设，所有的参数在每个分布式进程模型中注册的顺序一致。该 module
   自身将会按照该模型中参数注册的相反顺序执行梯度的
   all-reduction。换言之，用户应该保证，每个分布式进程模型一样，且参数注册顺序一致。(这也是由分布式原理决定)
6. 如果计划使用该 module，并用 NCCL 后端或 Gloo 后端 (使用 infiniband)，需要与多 workers 的
   Dataloader 一同使用，请修改多进程启动算法为 forkserver (python 3 only) 或 spawn
   。不幸的是，Gloo (使用 infiniband) 和 NCCL2 fork 并不安全，并且如果不改变配置时，很可能会
   deadlocks。
7. 在 module 上定义的前向传播和反向传播 hooks 和其子 modules 将不会涉及，除非 hooks 在 forward
   中进行了初始化。
8. 在使用 DistributedDataParallel 封装 model 后，不应该再修改模型的参数。也就是说，当使用
   DistributedDataParallel 打包 model 时，DistributedDataParallel 的
   constructor 将会在模型上注册额外的归约函数，该函数作用于模型的所有参数。

如果在构建 DistributedDataParallel 之后，改变模型的参数，这是不被允许的，并且可能会导致不可预期的后果，因为部分参数的梯度归约函数可能不会被调用。

在进程之间，参数永远不会进行 broadcast。该 module 对梯度执行一个 all-reduce 步骤，并假设在所有进程中，可以被 optimizer 以相同的方式进行更改。 在每一次迭代中，Buffers (BatchNorm stats 等) 是进行 broadcast 的，从 rank 0 的进程中的 module 进行广播，广播到系统的其他副本中。

1.4 torch.distribution.launch

该部分内容请看博主这篇文章——torch.distribution.launch

1.5 总结

多机多卡的分布式跟单机多卡分布式很像，最大的不同其实是参数使用的不同，上文都已经写到。在进行多机器分布式时，除了启动命令不同，源代码都得一样，否则很容易报错。整体这个系列都是使用pytorch框架做的测试。文章会有很多不足，请各位看官多多指教。

用一个实例来完成本篇文章：

"""
(MNMC) Multiple Nodes Multi-GPU Cards Training
    with DistributedDataParallel and torch.distributed.launch
Try to compare with [snsc.py, snmc_dp.py & mnmc_ddp_mp.py] and find out the differences.
"""

import os

import torch
import torch.distributed as dist
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.nn.parallel import DistributedDataParallel as DDP

BATCH_SIZE = 256
EPOCHS = 5


if __name__ == "__main__":

    # 0. set up distributed device
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(rank % torch.cuda.device_count())
    dist.init_process_group(backend="nccl")
    device = torch.device("cuda", local_rank)

    print(f"[init] == local rank: {local_rank}, global rank: {rank} ==")

    # 1. define network
    net = torchvision.models.resnet18(pretrained=False, num_classes=10)
    net = net.to(device)
    # DistributedDataParallel
    net = DDP(net, device_ids=[local_rank], output_device=local_rank)

    # 2. define dataloader
    trainset = torchvision.datasets.CIFAR10(
        root="./data",
        train=True,
        download=False,
        transform=transforms.Compose(
            [
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(
                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                ),
            ]
        ),
    )
    # DistributedSampler
    # we test single Machine with 2 GPUs
    # so the [batch size] for each process is 256 / 2 = 128
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        trainset,
        shuffle=True,
    )
    train_loader = torch.utils.data.DataLoader(
        trainset,
        batch_size=BATCH_SIZE,
        num_workers=4,
        pin_memory=True,
        sampler=train_sampler,
    )

    # 3. define loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        net.parameters(),
        lr=0.01 * 2,
        momentum=0.9,
        weight_decay=0.0001,
        nesterov=True,
    )

    if rank == 0:
        print("            =======  Training  ======= \n")

    # 4. start to train
    net.train()
    for ep in range(1, EPOCHS + 1):
        train_loss = correct = total = 0
        # set sampler
        train_loader.sampler.set_epoch(ep)

        for idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)

            loss = criterion(outputs, targets)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            total += targets.size(0)
            correct += torch.eq(outputs.argmax(dim=1), targets).sum().item()

            if rank == 0 and ((idx + 1) % 25 == 0 or (idx + 1) == len(train_loader)):
                print(
                    "   == step: [{:3}/{}] [{}/{}] | loss: {:.3f} | acc: {:6.3f}%".format(
                        idx + 1,
                        len(train_loader),
                        ep,
                        EPOCHS,
                        train_loss / (idx + 1),
                        100.0 * correct / total,
                    )
                )
    if rank == 0:
        print("\n            =======  Training Finished  ======= \n")

# node1
>>> python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --nnodes=2 \
    --node_rank=0 \
    --master_addr="10.198.189.10" \
    --master_port=22222 \
    mnmc_ddp_launch.py

>>> python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --nnodes=2 \
    --node_rank=1 \
    --master_addr="10.198.189.10" \
    --master_port=22222 \
    mnmc_ddp_launch.py

node 2
>>> python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --nnodes=2 \
    --node_rank=1 \
    --master_addr="10.198.189.10" \
    --master_port=22222 \
    mnmc_ddp_launch.py

结果：

[init] == local rank: 5, global rank: 5 ==
[init] == local rank: 3, global rank: 3 ==
[init] == local rank: 2, global rank: 2 ==
[init] == local rank: 4, global rank: 4 ==
[init] == local rank: 0, global rank: 0 ==
[init] == local rank: 6, global rank: 6 ==
[init] == local rank: 7, global rank: 7 ==
[init] == local rank: 1, global rank: 1 ==
            =======  Training  ======= 
   == step: [ 13/13] [0/5] | loss: 2.056 | acc: 23.776%
   == step: [ 13/13] [1/5] | loss: 1.688 | acc: 36.736%
   == step: [ 13/13] [2/5] | loss: 1.508 | acc: 44.544%
   == step: [ 13/13] [3/5] | loss: 1.462 | acc: 45.472%
   == step: [ 13/13] [4/5] | loss: 1.357 | acc: 49.344%
            =======  Training Finished  ======= 

资料：

• https://zhuanlan.zhihu.com/p/76638962
• https://blog.csdn.net/kejizuiqianfang/article/details/102454278
• https://www.jianshu.com/p/5f6cd6b50140
• https://github.com/BIGBALLON/distribuuuu
• https://github.com/tczhangzhi/pytorch-distributed