DDP及其在pytorch中应用

曾经终败给现在 2022-12-01 14:16 494阅读 0赞

# 1 分布式训练及其分类 #

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

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

**数据并行 vs 模型并行**：

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

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center]相比于模型并行，数据并行更为常用。DP(Data Parallel)和DDP(Distributed Data Parallel)都属于数据并行。

**同步更新 vs 异步更新**：

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

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

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 1]  
![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 2][]  
同步更新有等待，速度取决于最慢的那个GPU；异步更新没有等待，但是会出现loss异常抖动等问题，实践中，一般使用的是同步更新。

**PS算法 vs Ring AllReduce算法**：

*  PS算法，parameter server，GPU 0将数据分为五份发到各GPU上，各GPU用自己分到的数据计算梯度，返回给GPU 0进行梯度累加求均值，然后进行参数更新，再由GPU 0将更新后的参数发给各GPU；  
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 3]  
    ps算法类似于map-reduce，其主要有两个缺点：
    
     *  每一轮迭代都需要将所有GPU的数据都同步完成后才能完成，并行的卡数目很多时，木桶效应很严重，计算效率低；
     *  所有的GPU都需要和Reducer进行数据、参数和梯度的通信，当模型较大或者数据较多时，通信开销很大。假设有N个GPU，一次完整通信的耗时为K，则PS算法的通信总耗时为 T = 2 ( N − 1 ) K T = 2(N - 1)K T=2(N−1)K，耗时随着GPU数量的增加而线性增长。
 *  Ring AllReduce算法： Ring AllReduce算法把所有的GPU组成一个逻辑环，每个GPU都从它的左邻居接收数据，向它的右邻居发送数据。Ring AllReduce分为两个步骤：Scatter Reduce和All Gather，Scatter Reduce用于在不同的GPU间交换数据，All Gather则用于在不同的GPU间同步数据。
    
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 4]
    
    Scatter Reduce：为简单起见，让我们假设目标是对一个浮点数的大数组求和; 系统中有N个GPU，每个GPU都有一个相同大小的数组，并且在allreduce的末尾，每个GPU都应该有一个相同大小的数组，其中包含原始数组中数字的总和。首先，gpu将数组划分为N个更小的块(其中N是环中的gpu数)。  
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 5]相邻的GPU传递不同的参数，每次传递都是把当前GPU的数据累加到其右邻居上，同时接收其左邻居的输入数据累加到自身，在传递N-1次之后，可以得到每一份参数的累积（在不同的GPU上）。  
    ![在这里插入图片描述][20200827164921886.gif_pic_center]
    
    All Gather：得到每一份参数的累积之后，再做一次传递，将累加结果同步复制到所有的GPU上。
    
    ![在这里插入图片描述][20200827165300330.gif_pic_center]  
    All Reduce的通信成本为： T = 2 ( N − 1 ) K N T = 2(N-1)\\frac\{K\}\{N\} T=2(N−1)NK， K N \\frac\{K\}\{N\} NK是因为每次只传递了 1 N \\frac\{1\}\{N\} N1的数据。可以看出，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][https_www.cnblogs.com_zf-blog_p_12010742.html]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 6]  
针对负载不均衡的问题，有人提出了分布式进行损失和初步梯度计算的解决方案，如下图所示：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 7]  
个人感觉，DDP出现之后，因为其训练效率更高，DP的应用空间已经很少，所以重点介绍DDP模式。

## 2.2 DDP ##

本部分引用自：[https://zhuanlan.zhihu.com/p/178402798][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][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_www.zhihu.com_question_57799212_answer_612786337]  
[https://pytorch.org/tutorials/beginner/saving\_loading\_models.html\#saving-torch-nn-dataparallel-models][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_72939003]  
[https://zhuanlan.zhihu.com/p/74792767][https_zhuanlan.zhihu.com_p_74792767]  
[https://zhuanlan.zhihu.com/p/75318339][https_zhuanlan.zhihu.com_p_75318339]  
[https://zhuanlan.zhihu.com/p/76638962][https_zhuanlan.zhihu.com_p_76638962]  
[https://zhuanlan.zhihu.com/p/178402798][https_zhuanlan.zhihu.com_p_178402798]  
[https://zhuanlan.zhihu.com/p/187610959][https_zhuanlan.zhihu.com_p_187610959]

[https_zhuanlan.zhihu.com_p_72939003]: https://zhuanlan.zhihu.com/p/72939003
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center]: /images/20221123/02d87b1b0cea472ba6f005523038f23f.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 1]: /images/20221123/c9f73a7a9c7c4213adc66616ff26f805.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 2]: /images/20221123/c8b3aaea782f449790f1ddc5bd3acde9.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 3]: /images/20221123/8f0c7e6747394702954fabb9bd034cb9.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 4]: /images/20221123/bd84a75303a6431ca7578bcefec5f972.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 5]: /images/20221123/9314c2900b4e4db0b3763096ef8b611c.png
[20200827164921886.gif_pic_center]: /images/20221123/9bafe902cc5e4780989b644087a37b20.png
[20200827165300330.gif_pic_center]: /images/20221123/863fa7453e654861a3d7df776241b992.png
[https_www.cnblogs.com_zf-blog_p_12010742.html]: https://www.cnblogs.com/zf-blog/p/12010742.html
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 6]: /images/20221123/709f5b01494f4d21a3264a031b3c46c2.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 7]: /images/20221123/872888e867cd447ea8632028cd42424b.png
[https_zhuanlan.zhihu.com_p_178402798]: https://zhuanlan.zhihu.com/p/178402798
[https_pytorch.org_docs_stable_modules_torch_multiprocessing_spawn.html_spawn]: https://pytorch.org/docs/stable/_modules/torch/multiprocessing/spawn.html#spawn
[https_www.zhihu.com_question_57799212_answer_612786337]: https://www.zhihu.com/question/57799212/answer/612786337
[https_pytorch.org_tutorials_beginner_saving_loading_models.html_saving-torch-nn-dataparallel-models]: https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-torch-nn-dataparallel-models
[https_zhuanlan.zhihu.com_p_74792767]: https://zhuanlan.zhihu.com/p/74792767
[https_zhuanlan.zhihu.com_p_75318339]: https://zhuanlan.zhihu.com/p/75318339
[https_zhuanlan.zhihu.com_p_76638962]: https://zhuanlan.zhihu.com/p/76638962
[https_zhuanlan.zhihu.com_p_187610959]: https://zhuanlan.zhihu.com/p/187610959