DDP及其在pytorch中应用

1 分布式训练及其分类

本部分内容引用自：https://zhuanlan.zhihu.com/p/72939003
分布式训练是为了应用多卡显卡加速模型的训练，可以分为三类：

• 按照并行方式分，可分为模型并行和数据并行；
• 按照更新方式分，可分为同步更新和异步更新；
• 按照算法来分，可分为Parameter Server和Ring AllReduce算法。

数据并行 vs 模型并行：

• 模型并行：模型大到单个显卡放不下的地步，就把模型分为几个部分分别放于不同的显卡上单独运行，各显卡输入相同的数据；
• 数据并行：不同的显卡输入不同的数据，运行完全相同的完整的模型。

相比于模型并行，数据并行更为常用。DP(Data Parallel)和DDP(Distributed Data Parallel)都属于数据并行。

同步更新 vs 异步更新：

数据并行中，每个显卡通过自己拥有的那一份数据进行前向推理和梯度计算，根据梯度进行模型权重更新时，就涉及到了参数更新方式使用同步更新还是异步更新。

• 同步更新：所有的GPU都计算完梯度后，累加到一起求均值进行参数更新，再进行下一轮的计算；
• 异步更新：每个GPU计算完梯度后，无需等待其他更新，立即更新参数并同步。

同步更新有等待，速度取决于最慢的那个GPU；异步更新没有等待，但是会出现loss异常抖动等问题，实践中，一般使用的是同步更新。

PS算法 vs Ring AllReduce算法：

• PS算法，parameter server，GPU 0将数据分为五份发到各GPU上，各GPU用自己分到的数据计算梯度，返回给GPU 0进行梯度累加求均值，然后进行参数更新，再由GPU 0将更新后的参数发给各GPU；
  
  ps算法类似于map-reduce，其主要有两个缺点：

  • 每一轮迭代都需要将所有GPU的数据都同步完成后才能完成，并行的卡数目很多时，木桶效应很严重，计算效率低；
  • 所有的GPU都需要和Reducer进行数据、参数和梯度的通信，当模型较大或者数据较多时，通信开销很大。假设有N个GPU，一次完整通信的耗时为K，则PS算法的通信总耗时为$T=2(N-1)K$，耗时随着GPU数量的增加而线性增长。

• Ring AllReduce算法： Ring AllReduce算法把所有的GPU组成一个逻辑环，每个GPU都从它的左邻居接收数据，向它的右邻居发送数据。Ring AllReduce分为两个步骤：Scatter Reduce和All Gather，Scatter Reduce用于在不同的GPU间交换数据，All Gather则用于在不同的GPU间同步数据。

  

  Scatter Reduce：为简单起见，让我们假设目标是对一个浮点数的大数组求和; 系统中有N个GPU，每个GPU都有一个相同大小的数组，并且在allreduce的末尾，每个GPU都应该有一个相同大小的数组，其中包含原始数组中数字的总和。首先，gpu将数组划分为N个更小的块(其中N是环中的gpu数)。
  相邻的GPU传递不同的参数，每次传递都是把当前GPU的数据累加到其右邻居上，同时接收其左邻居的输入数据累加到自身，在传递N-1次之后，可以得到每一份参数的累积（在不同的GPU上）。
  

  All Gather：得到每一份参数的累积之后，再做一次传递，将累加结果同步复制到所有的GPU上。

  
  All Reduce的通信成本为：$T=2(N-1)\frac{K}{N}$，$\frac{K}{N}$是因为每次只传递了$\frac{1}{N}$的数据。可以看出，All Reduce的通信成本和GPU数量无关。

2 pytorch中的DP和DDP

2.1 DP

DP模式是很早就出现的、单机多卡的、参数服务器架构的多卡训练模式，在PyTorch，即是：

from torch.nn import DataParallel

device = torch.device("cuda")

model = MyModel()
model = model.to(device)
model = DataParallel(model)

在DP模式中，总共只有一个进程，多个线程（受到GIL很强限制）。master节点相当于参数服务器，其会向其他卡广播其参数；在梯度反向传播后，各卡将梯度集中到master节点，master节点对搜集来的参数进行平均后更新参数，再将参数统一发送到其他卡上。这种参数更新方式，会导致master节点的计算任务、通讯量很重，从而导致网络阻塞，降低了训练速度。

DP模型的另一个缺点就是master节点的负载很重，因为把损失和损失相对于最后一层的梯度计算都放在了master节点上。如下图所示：
下面两幅图参考自：https://www.cnblogs.com/zf-blog/p/12010742.html

针对负载不均衡的问题，有人提出了分布式进行损失和初步梯度计算的解决方案，如下图所示：

个人感觉，DDP出现之后，因为其训练效率更高，DP的应用空间已经很少，所以重点介绍DDP模式。

2.2 DDP

本部分引用自：https://zhuanlan.zhihu.com/p/178402798

DDP支持单机多卡分布式训练，也支持多机多卡分布式训练。目前DDP模式只能在Linux下应用。

DDP模式相对于DP模式的最大区别是启动了多个进程进行并行训练，用过python的人都了解，python代码运行的时候需要使用GIL进行解释，而每个进程只有一个GIL。因此，对比DDP和DP，可以明确DDP效率优于DP。

DDP有不同的使用模式。DDP的官方最佳实践是，每一张卡对应一个单独的GPU模型（也就是一个进程），在下面介绍中，都会默认遵循这个pattern。

理论上，DDP性能和单卡Gradient Accumulation性能是完全一致的。并行度为8的DDP 等于 Gradient Accumulation Step为8的单卡速度上，DDP当然比Graident Accumulation的单卡快；

2.2.1 pytorch中使用DDP

2.2.1.1 参数

在16张显卡，16个显卡的并行训练下，DDP会同时启动16个进程。

group

即进程组。默认情况下，只有一个组。

world size

表示全局的进程数，简单来讲，就是2x8=16。

	# 获取world size，在不同进程里都是一样的，得到16
	torch.distributed.get_world_size()

rank

当前进程的序号，用于进程间通讯。对于16的world size来说，就是0,1,2,…,15。
注意：rank=0的进程就是master进程。

	# 获取rank，每个进程都有自己的序号，各不相同
	torch.distributed.get_rank()

local_rank
又一个序号。这是每台机子上的进程的序号。机器一上有0,1,2,3,4,5,6,7，机器二上也有0,1,2,3,4,5,6,7。

	# 获取local_rank。一般情况下，你需要用这个local_rank来手动设置当前模型是跑在当前机器的哪块GPU上面的。
	torch.distributed.local_rank()

2.2.1.2 模型

模型转换成DDP模型

	## main.py文件
	import torch
	import argparse
	
	# 新增1:依赖
	import torch.distributed as dist
	from torch.nn.parallel import DistributedDataParallel as DDP
	
	# 新增2：从外面得到local_rank参数，在调用DDP的时候，其会自动给出这个参数，后面还会介绍。所以不用考虑太多，照着抄就是了。
	#       argparse是python的一个系统库，用来处理命令行调用，如果不熟悉，可以稍微百度一下，很简单！
	parser = argparse.ArgumentParser()
	parser.add_argument("--local_rank", default=-1)
	FLAGS = parser.parse_args()
	local_rank = FLAGS.local_rank
	
	# 新增3：DDP backend初始化
	#   a.根据local_rank来设定当前使用哪块GPU
	torch.cuda.set_device(local_rank)
	#   b.初始化DDP，使用默认backend(nccl)就行。如果是CPU模型运行，需要选择其他后端。
	dist.init_process_group(backend='nccl')
	
	# 新增4：定义并把模型放置到单独的GPU上，需要在调用`model=DDP(model)`前做哦。
	#       如果要加载模型，也必须在这里做哦。
	device = torch.device("cuda", local_rank)
	model = nn.Linear(10, 10).to(device)
	# 可能的load模型...
	
	# 新增5：之后才是初始化DDP模型
	model = DDP(model, device_ids=[local_rank], output_device=local_rank)

model = DDP(model, device_ids=[local_rank], output_device=local_rank)就是把原有的model变为DDP模式的model。

2.2.1.3 训练数据

我们知道，DDP同时起了很多个进程，但是他们用的是同一份数据，那么就会有数据上的冗余性。也就是说，你平时一个epoch如果是一万份数据，现在16个进程就要变成1*16=16万份数据了。

那么，我们需要使用一个特殊的sampler，来使得各个进程上的数据各不相同，进而让一个epoch还是1万份数据。也就是使用下面代码中的torch.distributed.DistributedSampler来创建分布式sampler。

	my_trainset = torchvision.datasets.CIFAR10(root='./data', train=True)
	# 新增1：使用DistributedSampler，DDP帮我们把细节都封装起来了。用，就完事儿！
	#       sampler的原理，后面也会介绍。
	train_sampler = torch.distributed.DistributedSampler(my_trainset)
	# 需要注意的是，这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
	trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)
	
	
	for epoch in range(num_epochs):
	    # 新增2：设置sampler的epoch，DistributedSampler需要这个来维持各个进程之间的相同随机数种子
	    trainloader.sampler.set_epoch(epoch)
	    # 后面这部分，则与原来完全一致了。
	    for data, label in trainloader:
	        prediction = model(data)
	        loss = loss_fn(prediction, label)
	        loss.backward()
	        optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
	        optimizer.step()

2.2.1.4 分布式模型的保存

save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。 因为model其实是DDP model，参数是被model=DDP(model)包起来的。

只需要在进程0上保存一次就行了，避免多次保存重复的东西。

	if dist.get_rank() == 0:
	    torch.save(model.module, "saved_model.ckpt")

2.2.1.5 DDP启动

我们需要在每一台机子（总共m台，每台n个显卡）上都运行一次torch.distributed.launch。每个torch.distributed.launch会启动n个进程，并给每个进程一个–local_rank = i 的参数。这就是之前需要"新增:从外面得到local_rank参数"的原因这样我们就得到n * m个进程，world_size=n * m。

1. 用torch.distributed.launch：

   单机模式：

   假设我们只在一台机器上运行，可用卡数是8。

   ## Bash运行
   python -m torch.distributed.launch --nproc_per_node=8 --batch-size=256 main.py
   

   • --nproc_per_node specifies how many GPUs you would like to use. In the example above, it is 8.
   • --batch-size is now the Total batch-size. It will be divided evenly to each GPU. In the example above, it is 256/2=32 per GPU.

   多机模式：
   多机模式需要两个额外的参数：

   • master_address：也就是master进程的网络地址默认是：127.0.0.1，只能用于单机。

   • 通讯的port：master_port 也就是master进程的一个端口，要先确认这个端口没有被其他程序占用了哦。一般情况下用默认的就行默认是：29500

     ## Bash运行
     # 假设我们在2台机器上运行，每台可用卡数是8
     #    机器1：
     python -m torch.distributed.launch --nnodes=2 --node_rank=0 --nproc_per_node 8 \
       --master_adderss $my_address --master_port $my_port main.py
     #    机器2：
     python -m torch.distributed.launch --nnodes=2 --node_rank=1 --nproc_per_node 8 \
       --master_adderss $my_address --master_port $my_port main.py
     

   如果只使用指定的显卡：

   	# 假设我们只用4,5,6,7号卡
   	CUDA_VISIBLE_DEVICES="4,5,6,7" python -m torch.distributed.launch --nproc_per_node 4 main.py
   	# 假如我们还有另外一个实验要跑，也就是同时跑两个不同实验。
   	#    这时，为避免master_port冲突，我们需要指定一个新的。这里我随便敲了一个。
   	CUDA_VISIBLE_DEVICES="4,5,6,7" python -m torch.distributed.launch --nproc_per_node 4 \
   	    --master_port 53453 main.py
   

2. 用mp.spawn

   PyTorch引入了torch.multiprocessing.spawn，可以使得单卡、DDP下的外部调用一致，即不用使用torch.distributed.launch。 python main.py一句话搞定DDP模式。

   文档：https://pytorch.org/docs/stable/_modules/torch/multiprocessing/spawn.html#spawn

   	def demo_fn(rank, world_size):
   	    dist.init_process_group("nccl", rank=rank, world_size=world_size)
   	    # lots of code.
   	    ...
   	
   	def run_demo(demo_fn, world_size):
   	    mp.spawn(demo_fn,
   	             args=(world_size,),
   	             nprocs=world_size,
   	             join=True)
   

2.2.2 DDP代码综合：

	################
	## main.py文件
	import torch
	# 新增：
	import torch.distributed as dist
	from torch.nn.parallel import DistributedDataParallel as DDP
	
	# 新增：从外面得到local_rank参数
	import argparse
	parser = argparse.ArgumentParser()
	parser.add_argument("--local_rank", default=-1)
	FLAGS = parser.parse_args()
	local_rank = FLAGS.local_rank
	
	# 新增：DDP backend初始化
	torch.cuda.set_device(local_rank)
	dist.init_process_group(backend='nccl')  # nccl是GPU设备上最快、最推荐的后端
	
	# 构造模型
	device = torch.device("cuda", local_rank)
	model = nn.Linear(10, 10).to(device)
	# load模型要在构造DDP模型之前，且只需要在master上加载就行了。
	if dist.get_rank() == 0:
	    model.load_state_dict(torch.load(ckpt_path))
	# 新增：构造DDP model
	model = DDP(model, device_ids=[local_rank], output_device=local_rank)
	
	# 优化器：要在构造DDP model之后，才能初始化model。
	optimizer = optim.SGD(model.parameters(), lr=0.001)
	
	# 构造数据
	my_trainset = torchvision.datasets.CIFAR10(root='./data', train=True)
	# 新增：使用DistributedSampler，DDP帮我们把细节都封装起来了。用，就完事儿！
	#       sampler的原理，后面也会介绍。
	train_sampler = torch.distributed.DistributedSampler(my_trainset)
	# 需要注意的是，这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
	trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, 
					sampler=train_sampler)
	
	# 网络训练
	model.train()
	for epoch in range(num_epochs):
	    # 新增：设置sampler的epoch，DistributedSampler需要这个来维持各个进程之间的相同随机数种子
	    trainloader.sampler.set_epoch(epoch)
	    # 后面这部分，则与原来完全一致了。
	    for data, label in trainloader:
	        optimizer.zero_grad()
	        prediction = model(data)
	        loss = loss_fn(prediction, label)
	        loss.backward()
	        optimizer.step()
	    # 1. save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。
	    #    因为model其实是DDP model，参数是被`model=DDP(model)`包起来的。
	    # 2. 只需要在进程0上保存一次就行了，避免多次保存重复的东西。
	    if dist.get_rank() == 0:
	        torch.save(model.module.state_dict(), "%d.ckpt" % epoch)

## Bash运行
# 改变：使用torch.distributed.launch启动DDP模式，
#   其会给main.py一个local_rank的参数。这就是之前需要"新增:从外面得到local_rank参数"的原因
python -m torch.distributed.launch --nproc_per_node 4 main.py

参考：https://www.zhihu.com/question/57799212/answer/612786337
https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-torch-nn-dataparallel-models
https://zhuanlan.zhihu.com/p/72939003
https://zhuanlan.zhihu.com/p/74792767
https://zhuanlan.zhihu.com/p/75318339
https://zhuanlan.zhihu.com/p/76638962
https://zhuanlan.zhihu.com/p/178402798
https://zhuanlan.zhihu.com/p/187610959