Pytorch 分布式训练（DP/DDP）

ytusdc

已于 2023-01-14 15:04:29 修改

阅读量2.9w

点赞数 104

分类专栏： Deep Learning 文章标签： pytorch 分布式深度学习

于 2021-12-23 14:35:50 首次发布

本文链接：https://blog.csdn.net/ytusdc/article/details/122091284

版权

Deep Learning 专栏收录该内容

33 篇文章 9 订阅

订阅专栏

参考（需要看）：

0、实操教程 | GPU多卡并行训练总结（以pytorch为例）

1、PyTorch 源码解读之 DP & DDP：模型并行和分布式训练解析 - 知乎

2、pytorch中分布式训练DP、DDP原理 - 知乎

3、pytorch中多卡训练_ytusdc的博客-CSDN博客_pytorch多卡训练

4、github代码机械：https://github.com/ytusdc/pytorch-distributed

5、PyTorch 分布式数据并行入门_w3cschool

6、Pytorch一机多卡分布式并行训练及混合精度训练_David's Tweet-CSDN博客_pytorch混合精度训练

7、深度学习分布式训练实战（一）_小肥柴不是小废柴的博客-CSDN博客_分布式深度学习

8、如何理解深度学习分布式训练中的large batch size与learning rate的关系？-51CTO.COM

9、Broadcast,Scatter,Gather,Reduce,All-reduce分别是什么？_aiwanghuan5017的博客-CSDN博客

10、Ring AllReduce简介 - 墨天轮

11、PyTorch分布式官方文档

12、官方文档写PyTorch分布式程序

1、分布式训练：模型并行和数据并行

分布式训练根据并行策略的不同，可以分为模型并行和数据并行。

模型并行

模型并行主要应用于模型相比显存来说更大，一块 GPU 无法加载的场景，通过把模型切割为几个部分，分别加载到不同的 GPU 上，来进行训练
数据并行

这个是日常会应用的比较多的情况。即每个 GPU 复制一份模型，将一批样本分为多份分发到各个GPU模型并行计算。因为求导以及加和都是线性的，数据并行在数学上也有效。采用数据并行相当于加大了batch_size，得到更准确的梯度或者加速训练。

模型并行不是这里的讨论重点

常用的 API 有两个：

torch.nn.DataParallel(DP)
torch.nn.DistributedDataParallel(DDP)

DP 相比 DDP 使用起来更友好（代码少），但是 DDP 支持多机多卡，训练速度更快，而且负载相对要均衡一些。所以优先选用 DDP 吧。

1.1、数据并行

数据并行的操作要求我们将数据划分成多份，然后发送给多个 GPU 进行并行的计算。

注意：多卡训练要考虑通信开销的，是个trade off的过程，不见得四块卡一定比两块卡快多少，可能是训练到四块卡的时候通信开销已经占了大头

下面是一个简单的示例。要实现数据并行，第一个方法是采用 nn.parallel 中的几个函数，分别实现的功能如下所示：

复制(Replicate)：将模型拷贝到多个 GPU 上；
分发(Scatter)：将输入数据根据其第一个维度(通常就是 batch 大小)划分多份，并传送到多个 GPU 上；
收集(Gather)：从多个 GPU 上传送回来的数据，再次连接回一起；
并行的应用(parallel_apply)：将第三步得到的分布式的输入数据应用到第一步中拷贝的多个模型上。

实现代码如下

# Replicate module to devices in device_ids
replicas = nn.parallel.replicate(module, device_ids)
# Distribute input to devices in device_ids
inputs = nn.parallel.scatter(input, device_ids)
# Apply the models to corresponding inputs
outputs = nn.parallel.parallel_apply(replicas, inputs)
# Gather result from all devices to output_device
result = nn.parallel.gather(outputs, output_device)

PyTorch也提供了简单的函数 torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)，只用几行代码可实现简单高效的并行GPU计算。通过device_ids参数可以指定在哪些GPU上进行优化，output_device指定输出到哪个GPU上

new_net= nn.DataParallel(net, device_ids=[0, 1])
output= new_net(input)

1.2、并行数据加载

流行的深度学习框架（例如Pytorch和Tensorflow）为分布式培训提供内置支持。从广义上讲，从磁盘读取输入数据开始，加载数据涉及四个步骤：

将数据从磁盘加载到主机
将数据从可分页内存传输到主机上的固定内存。请参阅此有关分页和固定的内存更多信息。
将数据从固定内存传输到GPU
在GPU上向前和向后传递

PyTorch 中的 Dataloader 提供使用多个进程（通过将 num_workers> 0 设置）从磁盘加载数据以及将多页数据从可分页内存到固定内存的能力（通过设置 pin_memory = True）。

一般的，对于大批量的数据，若仅有一个线程用于加载数据，则数据加载时间占主导地位，这意味着无论我们如何加快数据处理速度，性能都会受到数据加载时间的限制。现在，设置num_workers = 4 以及 pin_memory = True。这样，可以使用多个进程从磁盘读取不重叠的数据，并启动生产者-消费者线程以将这些进程读取的数据从可分页的内存转移到固定的内存。

多个进程能够更快地加载数据，并且当数据处理时间足够长时，流水线数据加载几乎可以完全隐藏数据加载延迟。这是因为在处理当前批次的同时，将从磁盘读取下一个批次的数据，并将其传输到固定内存。如果处理当前批次的时间足够长，则下一个批次的数据将立即可用。这个想法需要为num_workers 参数设置适当的值。设置此参数，以使从磁盘读取批处理数据的速度比GPU处理当前批处理的速度更快（但不能更高，因为这只会浪费多个进程使用的系统资源）。

请注意，到目前为止，我们仅解决了从磁盘加载数据以及从可分页到固定内存的数据传输问题。从固定内存到GPU的数据传输（tensor.cuda()）也可以使用CUDA流进行流水线处理。

现在将使用GPU网络检查数据并行处理。基本思想是，网络中的每个GPU使用模型的本地副本对一批数据进行正向和反向传播。反向传播期间计算出的梯度将发送到服务器，该服务器运行reduce归约操作以计算平均梯度。然后将平均梯度结果发送回GPU，GPU使用SGD更新模型参数。使用数据并行性和有效的网络通信软件库（例如NCCL），可以实现使训练时间几乎线性减少。

2、数据并行 DataParallel

PyTorch 中实现数据并行的操作可以通过使用 torch.nn.DataParallel

`2.1 并行处理机制`

DataParallel系统通过将整个小型批处理加载到主线程上，然后将子小型批处理分散到整个GPU网络中来工作。

详细流程：

forward：是将输入一个 batch 的数据均分成多份，分别送到对应的 GPU 进行计算。与 Module 相关的所有数据也都会以浅复制的方式复制多份。每个 GPU 在单独的线程上将针对各自的输入数据独立并行地进行 forward 计算。

backward：在主GPU上收集网络输出，并通过将网络输出与批次中每个元素的真实数据标签进行比较来计算损失函数值。接下来，损失值分散给各个GPU，每个GPU进行反向传播以计算梯度。最后，在主GPU上归约梯度、进行梯度下降，并更新主GPU上的模型参数。由于模型参数仅在主GPU上更新，而其他从属GPU此时并不是同步更新的，所以需要将更新后的模型参数复制到剩余的从属 GPU 中，以此来实现并行。

DataParallel会将定义的网络模型参数默认放在GPU 0上，所以dataparallel实质是可以看做把训练参数从GPU拷贝到其他的GPU同时训练，这样会导致内存和GPU使用率出现很严重的负载不均衡现象，即GPU 0的使用内存和使用率会大大超出其他显卡的使用内存，因为在这里GPU0作为master来进行梯度的汇总和模型的更新，再将计算任务下发给其他GPU，所以他的内存和使用率会比其他的高。

具体流程见下图：

官方多GPU下的forward和backward：

跟上述流程图稍微不一样，一个loss和Grad，

上图清晰的表达了 torch.nn.DataParallel 的计算过程。

将 inputs 从主 GPU 分发到所有 GPU 上
将 model 从主 GPU 分发到所有 GPU 上
每个 GPU 分别独立进行前向传播，得到 outputs
将每个 GPU 的 outputs 发回主 GPU
在主 GPU 上，通过 loss function 计算出 loss，对 loss function 求导，求出损失梯度
计算得到的梯度分发到所有 GPU 上
反向传播计算参数梯度
将所有梯度回传到主 GPU，通过梯度更新模型权重
不断重复上面的过程

总结流程：

基本上，给定的输入通过在批处理维度中分块在GPU之间进行分配。在前向传递中，模型在每个设备上复制，每个副本处理批次的一部分。在向后传递过程中，主GPU（上图中的GPU-1）收集每个GPU的输出output，根据label计算loss，继而计算得到多个梯度grad，然后将梯度分发到各个GPU（官方原理图中第二行第二个），然后每个GPU副本模型上的梯度更新（第二行第三个），然后再将每个更新完梯度的的参数合并到主gpu(第二行最后一个步骤)，求和以生成最终的梯度，并将其应用于主gpu（上图中的GPU-1）以更新模型权重。在下一次迭代中，主GPU上的更新模型将再次复制到每个GPU设备上。

优点：并行化多个GPU上的NN训练，因此与累积梯度相比，它减少了训练时间。因为代码更改很少，所以适合快速原型制作。

缺点：nn.DataParallel使用单进程多线程方法在不同的GPU上训练相同的模型。它将主进程保留在一个GPU上，并在其他GPU上运行不同的线程。由于python中的线程存在GIL（全局解释器锁定）问题，因此这限制了完全并行的分布式训练设置。

许多低效率之处：

冗余数据副本
- 数据从主机复制到主GPU，然后将子微型批分散在其他GPU上
在前向传播之前跨GPU进行模型复制
- 由于模型参数是在主GPU上更新的，因此模型必须在每次正向传递的开始时重新同步
每批的线程创建/销毁开销
- 并行转发是在多个线程中实现的（这可能只是PyTorch问题）
梯度减少流水线机会未开发
- 在Pytorch 1.0数据并行实现中，梯度下降发生在反向传播的末尾。
在主GPU上不必要地收集模型输出output
GPU利用率不均
- 在主GPU上执行损失loss计算
- 梯度下降，在主GPU上更新参数

2.2 代码示例

如果不设定好要使用的 device_ids 的话, 程序会自动找到这个机器上面可以用的所有的显卡用于训练。

如果想要限制使用的显卡数，怎么办呢？

在代码中加入最前面使用： os.environ['CUDA_VISIBLE_DEVICES'] == '0,5'

注意：这行代码要加在 import os之后，否则可能不起作用。例如

import os 
os.environ['CUDA_VISIBLE_DEVICES'] == '0,5'
import torch
import xxxx

限制代码能看到的GPU个数，这里表示指定只使用实际的0号和5号GPU # 注意：这里的赋值必须是字符串，list会报错

或者终端执行程序时 CUDA_VISIBLE_DEVICES=1,2,3 python main.py

如果不指定，nn.DataParallel会默认用能见到的所有GPU

#device_ids = [0,1,2,3]

if args.cuda:
    model =  model.cuda()    #这里将模型复制到gpu
if len(device_ids)>1:
  model = nn.DataParallel(model)

# 或者指定GPU
model = nn.DataParallel(model.cuda(), device_ids=[gpus])

class torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

- module ：待进行并行的模型

- device_ids : GPU 列表，其值可以是 torch.device 类型，也可以是 int list. 默认使用全部 GPUs.

- output_device : GPU ID 或 torch.device. 指定输出的 GPU，默认为第一个，即 device_ids[0].

2.3 cuda()函数解释

cuda() 函数返回一个存储在CUDA内存中的复制，其中device可以指定cuda设备。但如果此storage对象早已在CUDA内存中存储，并且其所在的设备编号与cuda()函数传入的device参数一致，则不会发生复制操作，返回原对象。

cuda()函数的参数信息:

device (int) – 指定的GPU设备id. 默认为当前设备，即 torch.cuda.current_device()的返回值。
non_blocking (bool) – 如果此参数被设置为True, 并且此对象的资源存储在固定内存上(pinned memory)，那么此cuda()函数产生的复制将与host端的原storage对象保持同步。否则此参数不起作用。

3 分布式数据并行 DistributedDataParallel

3.1 并行处理机制

DistributedDataParallel，支持 all-reduce，broadcast，send 和 receive 等等。通过 MPI 实现 CPU 通信，通过 NCCL 实现 GPU 通信。可以用于单机多卡也可用于多机多卡, 官方也曾经提到用 DistributedDataParallel 解决 DataParallel 速度慢，GPU 负载不均衡的问题。

效果比DataParallel好太多！！！torch.distributed 相对于 torch.nn.DataParalle 是一个底层的API，所以我们要修改我们的代码，使其能够独立的在机器（节点）中运行。

与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，我们只需要编写一份代码，torch 就会自动将其分配给n个进程，分别在 n 个 GPU 上运行。不再有主GPU，每个GPU执行相同的任务。对每个GPU的训练都是在自己的过程中进行的。每个进程都从磁盘加载其自己的数据。分布式数据采样器可确保加载的数据在各个进程之间不重叠。损失函数的前向传播和计算在每个GPU上独立执行。因此，不需要收集网络输出。在反向传播期间，梯度下降在所有GPU上均被执行，从而确保每个GPU在反向传播结束时最终得到平均梯度的相同副本。

3.2 对比DataParallel，DistributedDataParallel的区别和优势如下

区别：DDP通过多进程实现的。也就是说操作系统会为每个GPU创建一个进程,从而避免了Python解释器GIL带来的性能开销。而DataParallel()是通过单进程控制多线程来实现的。

优势：

1、每个进程对应一个独立的训练过程，且只对梯度等少量数据进行信息交换。

在每次迭代中，每个进程具有自己的 optimizer ，并独立完成所有的优化步骤，进程内与一般的训练无异。

在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。之后，各进程用该梯度来独立的更新参数。（而 DataParallel是梯度汇总到gpu0，反向传播更新参数，再广播模型参数给其他的gpu）。DDP由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致。

而在 DataParallel 中，全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。

相较于 DataParallel，torch.distributed 传输的数据量更少，因此速度更快，效率更高。

2、每个进程包含独立的解释器和 GIL。

一般使用的Python解释器CPython：是用C语言实现Pyhon，是目前应用最广泛的解释器。全局锁使Python在多线程效能上表现不佳，全局解释器锁（Global Interpreter Lock）是Python用于同步线程的工具，使得任何时刻仅有一个线程在执行。

由于每个进程拥有独立的解释器和 GIL，消除了来自单个 Python 进程中的多个执行线程，模型副本或 GPU 的额外解释器开销和 GIL-thrashing ，因此可以减少解释器和 GIL 使用冲突。这对于严重依赖 Python runtime 的 models 而言，比如说包含 RNN 层或大量小组件的 models 而言，这尤为重要。

3、为什么尽管增加了复杂性，但还是考虑使用DistributedDataParallel而不是DataParallel：

如果模型太大而无法容纳在单个 GPU 上，则必须使用模型并行将其拆分到多个 GPU 中。 DistributedDataParallel与模型并行一起使用； DataParallel目前没有。
DataParallel是单进程，多线程，并且只能在单台机器上运行，而DistributedDataParallel是多进程，并且适用于单机和多机训练。因此，即使在单机训练中，数据足够小以适合单机，DistributedDataParallel仍比DataParallel快。 DistributedDataParallel还预先复制模型，而不是在每次迭代时复制模型，并避免了全局解释器锁定。
如果您的两个数据都太大而无法容纳在一台计算机和上，而您的模型又太大了以至于无法安装在单个 GPU 上，则可以将模型并行(跨多个 GPU 拆分单个模型）与DistributedDataParallel结合使用。在这种情况下，每个DistributedDataParallel进程都可以并行使用模型，而所有进程都将并行使用数据。

4 分布式训练介绍

分布式训练可以分为：

单机多卡，DataParallel（最常用，最简单）-- 看前文
单机多卡，DistributedDataParallel（较高级）
多机多卡，DistributedDataParallel（最高级）

Pytorch分布训练一开始接触的往往是DataParallel，这个wrapper能够很方便的使用多张卡，而且将进程控制在一个。唯一的问题就在于，DataParallel只能满足一台机器上gpu的通信，而一台机器一般只能装8张卡，对于一些大任务，8张卡就很吃力了，这个时候我们就需要面对多机多卡分布式训练这个问题了

DistributedDataParallel (DDP）在模块级别实现数据并行性。它使用 Torch.distributed 程序包中的通信集合来同步梯度，参数和缓冲区。并行性在流程内和跨流程均可用。在一个过程中，DDP 将输入模块复制到device_ids中指定的设备，将输入沿批次维度分散，然后将输出收集到output_device，这与 DataParallel 相似。在整个过程中，DDP 在正向传递中插入必要的参数同步，在反向传递中插入梯度同步。用户可以将进程映射到可用资源，只要进程不共享 GPU 设备即可。推荐的方法(通常是最快的方法）是为每个模块副本创建一个过程，即在一个过程中不进行任何模块复制。

使用 nn.DistributedDataParallel 进行Multiprocessing可以在多个gpu之间复制该模型，每个gpu由一个进程控制。（如果你想，也可以一个进程控制多个GPU，但这会比控制一个慢得多。也有可能有多个工作进程为每个GPU获取数据，但为了简单起见，本文将省略这一点。）这些GPU可以位于同一个节点上，也可以分布在多个节点上。每个进程都执行相同的任务，并且每个进程与所有其他进程通信。只有梯度会在进程/GPU之间传播，这样网络通信就不至于成为一个瓶颈了。

训练过程中，每个进程从磁盘加载自己的小批（minibatch）数据，并将它们传递给自己的GPU。每个GPU都做它自己的前向计算，然后梯度在GPU之间全部约简。每个层的梯度不仅仅依赖于前一层，因此梯度全约简与并行计算反向传播，进一步缓解网络瓶颈。在反向传播结束时，每个节点都有平均的梯度，确保模型权值保持同步（synchronized）。

4.1 Pytorch 分布式使用流程

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

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

4.2 分布式几个基本概念

group： 即进程组。默认情况下，只有一个组，这个可以先不管，一直用默认的就行。一个 job 即为一个组，也即一个 world。

当需要进行更加精细的通信时，可以通过 new_group 接口，使用 word 的子集，创建新组，用于集体通信等。

world_size ：表示全局并行进程个数。

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

rank：表现当前进程的序号，用于进程间通讯，表征进程优先级。rank = 0 的主机为 master 节点。同时，rank=0的进程就是master进程

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

local_rank：这是每台机子上的进程的序号，进程内，GPU 编号，非显式参数，由 torch.distributed.launch 内部指定。比方说， rank = 3，local_rank = 0 表示第 3 个进程内的第 1 块 GPU。

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

举个栗子：4台机器(每台机器8张卡)进行分布式训练
通过 init_process_group() 对进程组进行初始化
初始化后可以通过 get_world_size() 获取到 world_size
在该例中为32，即有32个进程，其编号为0-31,即rank的取值范围，可以通过 get_rank() 函数可以进行获取。在每台机器上，local rank均为0-8，这是 local_rank 与 rank 的区别， local_rank 会对应到实际的 GPU ID 上
(单机多任务的情况下注意CUDA_VISIBLE_DEVICES的使用
控制不同程序可见的GPU devices)

4.3 几个基本函数介绍

1、torch.utils.data.distributed.DistributedSampler 的使用

在多机多卡情况下分布式训练数据的读取也是一个问题，不同的卡读取到的数据应该是不同的。dataparallel的做法是直接将batch切分到不同的卡，这种方法对于多机来说不可取，因为多机之间直接进行数据传输会严重影响效率。于是有了利用sampler确保dataloader只会load到整个数据集的一个特定子集的做法。DistributedSampler就是做这件事的。它为每一个子进程划分出一部分数据集，以避免不同进程之间数据重复

原型

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

参数

dataset 进行采样的数据集
num_replicas 分布式训练中，参与训练的进程数
rank 当前进程的 rank 序号（必须位于分布式训练中）

说明：

对数据集进行采样，使之划分为几个子集。

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

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

注意：在 DataParallel 中，batch size 设置必须为单卡的 n 倍，因为一个batch的数据会被主GPU分散为minibatch给其他GPU，但是在 DistributedDataParallel 内，batch size 设置于单卡一样即可，因为各个GPU对应的进程独立从磁盘中加载数据。

4.4 分布式的基本流程

4.4.1、添加参数 --local_rank

通过torch.distributed.launch来启动训练，torch.distributed.launch 会给每个进程分配一个 local_rank 参数，表示当前进程在当前主机上的编号。例如：rank=2, local_rank=0 表示第 3 个节点上的第 1 个进程。

local_rank代表当前程序进程使用的GPU标号
这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，也可以通过torch.distributed.get_rank()获取进程id，两种方法都可以

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()

# 通过args接收 local_rank
local_rank = args.local_rank

# 通过 get_rank() 得到 local_rank，必须放在初始化之后使用，否则会报错
#local_rank = torch.distributed.get_rank()

4.4.2、初始化使用nccl后端

# ps 检查nccl是否可用
# torch.distributed.is_nccl_available ()

torch.distributed.init_process_group(backend="nccl")

4.4.3、配置每个进程的gpu
local_rank代表当前程序进程使用的GPU标号

# 获取 local_rank
local_rank = torch.distributed.get_rank()
#配置每个进程的 GPU， 根据local_rank来设定当前使用哪块GPU
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

4.4.4、使用DistributedSampler

别忘了设置pin_memory=true
使用 DistributedSampler 对数据集进行划分。它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练
需要注意的是，这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
注意 testset不用sampler

# 自己的数据获取
dataset = MyDataset(input_size, data_size)

# 使用 DistributedSampler
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)

trainloader = DataLoader(dataset=dataset,
                         pin_memory=true,
                         shuffle=(train_sampler is None),   # 使用分布式训练 shuffle 应该设置为 False
                         batch_size=args.batch_size,
                         num_workers=args.workers,
                         sampler=train_sampler)

解释一下DataLoader中其中两个参数：

shuffle：

为了能够按顺序划分数据子集，拿到不同部分数据，所以数据集不能够进行随机打散，所以用了参数 'shuffle': False

num_worker：
数据集加载的时候，控制用于同时加载数据的线程数（默认为0，即在主线程读取）存在最优值，你会看到运行的时候pytorch会新建恰等于这个值的数据读取线程，我猜，线程多于必要的时候，数据读取线程返回到主线程反而会因为线程间通信减慢数据。因此大了不好小了也不好。建议把模型，loss，优化器全注释了只跑一下数据流速度，确定最优值
pin_memory：
是否提前申请CUDA内存（默认为False，但有说法除非数据集很小，否则在N卡上推荐总是打开）在MNIST这样的小数据集上好像是关闭比较好，到底多小算小说不清楚，建议自己试一下。

如果机子的内存比较大，建议开启pin_memory=Ture，如果开启后发现有卡顿现象或者内存占用过高，此时建议关闭。

为什么设置 pip_memory=true，看解释

多GPU训练的时候注意机器的内存是否足够(一般为使用显卡显存x2)，如果不够，建议关闭pin_memory(锁页内存)选项。
采用DistributedDataParallel多GPUs训练的方式比DataParallel更快一些，如果你的Pytorch编译时有nccl的支持，那么最好使用DistributedDataParallel方式。
关于什么是锁页内存：
pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
主机中的内存，有两种存在方式，一是锁页，二是不锁页，锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。显卡中的显存全部是锁页内存,当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。

4.4.5、分布式模型部署，引入SyncBN，将普通BN替换成SyncBN。

为什么使用 SyncBN看这里： [原创][深度][PyTorch] DDP系列第三篇：实战与技巧 - 知乎

使用 DistributedDataParallel 包装模型，它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值

BatchNorm之类的层在其计算中使用了整个批次统计信息，因此无法仅使用一部分批次在每个GPU上独立进行操作。 PyTorch提供SyncBatchNorm作为BatchNorm的替换/包装模块，该模块使用跨GPU划分的整个批次计算批次统计信息。

model = Model()
# 把模型移到对应的gpu
# 定义并把模型放置到单独的GPU上，需要在调用`model=DDP(model)`前做
model.to(device)

# 引入SyncBN，这句代码，会将普通BN替换成SyncBN。
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

# GPU 数目大于 1 才有必要分布式训练
if torch.cuda.device_count() > 1:
    model = torch.nn.parallel.DistributedDataParallel(model,
                                                      device_ids=[local_rank],
                                                      output_device=local_rank)

4.4.6、把数据和模型加载到当前进程使用的 GPU 中，正常进行正反向传播

这里要注意设置sampler的epoch，DistributedSampler需要这个来维持各个进程之间的相同随机数种子

for epoch in range(num_epochs):
    # 设置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()

4.5 并行程序启动

以上做完之后就是怎么启动程序，这里总共有有两种方法：

以下只是展示单机多卡情况，多机多卡在后面会有说明

4.5.1. 用 `torch.distributed.launch`：

torch.distributed.launch为我们触发了n个train.py 的 GPU进程(PID)，n就是我们将要使用的GPU或者需要开启的进程数数量。train.py会并行地n个运行。torch.distributed.launch 会给模型分配一个args.local_rank的参数，也可以通过torch.distributed.get_rank()获取进程id。

python -m torch.distributed.launch --nproc_per_node=4 main.py --{args}

4.5.2. 用 torch.multiprocessing：

注意：main函数的第一个 rank 变量会被 mp.spawn 函数自动填充，可以充当 local_rank 来用

import torch.multiprocessing as mp
 
def main(rank, your_custom_arg_1, your_custom_arg_2):
    pass # 将前面那一堆东西包装成一个 main 函数
 
mp.spawn(main, nprocs=how_many_process, args=(your_custom_arg_1, your_custom_arg_2))

以上步骤整理后的代码

import torch
import torch.nn as nn
import argparse
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
import os
from torch.utils.data.distributed import DistributedSampler

# 1)添加参数  --local_rank
'''
每个进程分配一个 local_rank 参数, 表示当前进程在当前主机上的编号。这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，local_rank代表当前程序进程使用的GPU标号
'''
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()

# 通过args接收 local_rank
local_rank = args.local_rank

# 通过 get_rank() 得到 local_rank，最好放在初始化之后使用
local_rank = torch.distributed.get_rank()

# 2)********************************************************
# 初始化使用nccl后端
torch.distributed.init_process_group(backend="nccl")

# 3)********************************************************
# 获取 local_rank
local_rank = torch.distributed.get_rank()
#配置每个进程的 GPU， 根据local_rank来设定当前使用哪块GPU
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

# 4)********************************************************
# 自己的数据获取
dataset = MyDataset(input_size, data_size)

# 使用 DistributedSampler
'''
使用 DistributedSampler 对数据集进行划分。它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练
#注意 testset不用sampler
'''
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)

trainloader = DataLoader(dataset=dataset,
                         pin_memory=true,
                         shuffle=(train_sampler is None),   # 使用分布式训练 shuffle 应该设置为 False
                         batch_size=args.batch_size,
                         num_workers=args.workers,
                         sampler=train_sampler)



# 5)********************************************************
model = Model()
# 把模型移到对应的gpu
# 定义并把模型放置到单独的GPU上，需要在调用`model=DDP(model)`前做
model.to(device)

# 引入SyncBN，这句代码，会将普通BN替换成SyncBN。
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

# GPU 数目大于 1 才有必要分布式训练
if torch.cuda.device_count() > 1:
    model = torch.nn.parallel.DistributedDataParallel(model,
                                                      device_ids=[local_rank],
                                                      output_device=local_rank)


# 6)********************************************************

for epoch in range(num_epochs):
    # 设置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()

在使用时，命令行调用 torch.distributed.launch 启动器启动：
pytorch 为我们提供了 torch.distributed.launch 启动器，用于在命令行分布式地执行 python 文件。
--nproc_per_node参数指定为当前主机创建的进程数。一般设定为=NUM_GPUS_YOU_HAVE当前主机的 GPU 数量，每个进程独立执行训练脚本。
这里是单机多卡，所以node=1，就是一台主机，一台主机上--nproc_per_node个进程

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

5、分布式初始化函数介绍

5.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, group_name='')

参数说明：

backend(str): 后端选择，包括上面那几种 gloo,nccl,mpi
init_method(str，optional): 用来初始化包的URL, 我理解是一个用来做并发控制的共享方式
world_size(int, optional): 参与这个工作的进程数
rank(int,optional): 当前进程的rank
group_name(str,optional): 用来标记这组进程名的

需要注意的是：

后端最好用“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)

torch.distributed初始化方式

torch.distributed.init_process_group(backend='nccl',
                                     init_method='env://')

5.2 torch.distributed.launch 启动程序

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

这个工具可以用作CPU或者GPU，如果被用于GPU，每个GPU产生一个进程Process

该工具既可以用来做单节点多GPU训练，也可用于多节点多GPU训练。如果是单节点多GPU，将会在单个GPU上运行一个分布式进程，据称可以非常好地改进单节点训练性能。如果用于多节点分布式训练，则通过在每个节点上产生多个进程来获得更好的多节点分布式训练性能。如果有Infiniband接口则加速比会更高。

在单节点分布式训练或多节点分布式训练的两种情况下，该工具将为每个节点启动给定数量的进程（--nproc_per_node）。如果用于GPU训练，则此数字需要小于或等于当前系统上的GPU数量（nproc_per_node），并且每个进程将在从GPU 0到GPU（nproc_per_node - 1）的单个GPU上运行。

单机多卡情况：Single-Node multi-process distributed training

在使用时，命令行调用 torch.distributed.launch 启动器启动：

pytorch 为我们提供了 torch.distributed.launch 启动器，用于在命令行分布式地执行 python 文件。

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py --{args}

说明：

CUDA_VISIBLE_DEVICES ：程序可用的GPU id号

--nproc_per_node ：参数指定为当前主机创建的进程数，一般设定为当前主机可用的GPU数目，即一个进行使用一个GPU，每个进程独立执行训练脚本
这里是单机多卡，所以node=1，就是总共一台主机，一台主机上--nproc_per_node个进程

多机多卡情况：Multi-Node multi-process distributed training

假设是2机3卡, nnode=2, 就是两台主机, 一台主机上--nproc_per_node=3个进程，命令应该如下：

在主机 01 上执行：

python torch.distributed.launch --nprocs_per_node=3 --nnodes=2 --node_rank=0 --master_addr="master-ip" --master_port=8888 main.py --my arguments

在主机 02 上执行：

python torch.distributed.launch --nprocs_per_node=3 --nnodes=2 --node_rank=1 --master_addr="master-ip" --master_port=8888  main.py --my arguments

参数说明：

nprocs_per_node：每台主机上要跑的进程数

nnodes ：总共几台主机，这里是 2 台主机

node_rank：当前是哪台主机，如上第一台主机上 node_rank=0，第二台主机上 node_rank=1

master_addr：必须设置，表示master 主机的IP地址，这里是第 1 台主机

master_port：必须设置，表示master 主机（rank = 0）的端口

但是当 python torch.distributed.launch 传递参数--use_env，其中的一些参数可以通过环境变量获得，具体没有测试过，例如一下可以获得

os.environ["RANK"]

os.environ['LOCAL_RANK']

os.environ['WORLD_SIZE']

5.3 使用 torch.multiprocessing 取代启动器

使用时，只需要调用 torch.multiprocessing.spawn，torch.multiprocessing 就会帮助我们自动创建进程。如下面的代码所示，spawn 开启了 nprocs=4 个进程，每个进程执行 main_worker 并向其中传入 local_rank（当前进程 index）和 args（即 4 和 myargs）作为参数：

import torch.multiprocessing as mp

mp.spawn(main_worker, nprocs=4, args=(4, myargs))

这里，我们直接将原本需要 torch.distributed.launch 管理的执行内容，封装进 main_worker 函数中，其中 proc 对应 local_rank（当前进程 index），进程数 nproc 对应 4， args 对应 myargs：

def main_worker(proc, nproc, args):

   dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)
   torch.cuda.set_device(args.local_rank)

   train_dataset = ...
   train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

   train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

   model = ...
   model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

   optimizer = optim.SGD(model.parameters())

   for epoch in range(100):
      for batch_idx, (data, target) in enumerate(train_loader):
          images = images.cuda(non_blocking=True)
          target = target.cuda(non_blocking=True)
          ...
          output = model(images)
          loss = criterion(output, target)
          ...
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()

在上面的代码中值得注意的是，由于没有 torch.distributed.launch 读取的默认环境变量作为配置，我们需要手动为 init_process_group 指定参数：

dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)

在使用时，直接使用 python 运行就可以了：

python main.py

6、分布式训练模型评估

在评估模型或生成日志时，需要从所有GPU收集当前批次统计信息，例如损失，准确率等，并将它们在一台机器上进行整理以进行日志记录。 PyTorch提供了以下方法，用于在所有GPU之间同步变量。

有关参数和方法的更多详细信息，请阅读torch.distributed软件包。 Distributed communication package - torch.distributed — PyTorch 1.10.1 documentation

torch.distributed.gather（input_tensor，collect_list，dst）：从所有设备收集指定的input_tensor并将它们放置在collect_list中的dst设备上。

torch.distributed.all_gather（tensor_list，input_tensor）：从所有设备收集指定的input_tensor并将其放置在所有设备上的tensor_list变量中。

torch.distributed.reduce（input_tensor，dst，reduce_op = ReduceOp.SUM）：收集所有设备的input_tensor并使用指定的reduce操作（例如求和，均值等）进行缩减。最终结果放置在dst设备上。

torch.distributed.all_reduce（input_tensor，reduce_op = ReduceOp.SUM）：与reduce操作相同，但最终结果被复制到所有设备。

示意图：

官方解释：Writing Distributed Applications with PyTorch — PyTorch Tutorials 1.10.1+cu102 documentation

参考文章：
PyTorch分布式训练详解教程 scatter, gather & isend, irecv & all_reduce & DDP - 天靖居士 - 博客园https://www.cnblogs.com/kkyyhh96/p/13769220.html torch.distributed【附源码】_wx5ba0c87f1984b_51CTO博客

分布式通信包 - torch.distributed - 简书

7、分布式训练保存模型

更多DDP模型保存和加载看另一篇文章这里：

pytorch GPU和CPU模型相互加载_ytusdc的博客-CSDN博客

需要注意的地方：

1. save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。因为model其实是DDP model，参数是被`model=DDP(model)`包起来的。
2. 只需要在进程0上，保存一次就行了，避免多次保存重复的东西。

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

２、.cuda或者.to(device)等问题

device是自己设置，如果.cuda出错，就要化成相应的device

model（如：model.to(device)）
input（通常需要使用Variable包装，如：input = Variable(input).to(device)）
target（通常需要使用Variable包装）
nn.CrossEntropyLoss()（如：criterion = nn.CrossEntropyLoss().to(device)）

其他参考文章：

Pytorch 分布式训练 - 知乎

PyTorch 21.单机多卡操作(分布式DataParallel，混合精度，Horovod) - 知乎

Pytorch中的Distributed Data Parallel与混合精度训练（Apex） - 知乎

pytorch(分布式)数据并行个人实践总结——DataParallel/DistributedDataParallel - fnangle - 博客园

[深度学习] 分布式Pytorch介绍（三）_小墨鱼的专栏-CSDN博客_world_size