Pytorch 分布式训练 (DP, DDP)

连理o

已于 2024-05-11 10:40:29 修改

阅读量7.6k

点赞数 15

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

于 2022-10-30 22:21:45 首次发布

本文链接：https://blog.csdn.net/weixin_42437114/article/details/127590671

版权

PyTorch 专栏收录该内容

10 篇文章 3 订阅

订阅专栏

`nn.DataParallel()` (DP)
`nn.parallel.DistributedDataParallel` (DDP)
其他问题
- 关于 batch_size 参数和 real batch size
References

`nn.DataParallel()` (DP)

多卡训练原理

基本概念

broadcast 是主进程将相同的数据分发给组里的每一个其它进程
scatter 是主进程将数据的每一小部分给组里的其它进程
gather 是将其它进程的数据收集过来；reduce 是将其它进程的数据收集过来并应用某种操作 (e.g. SUM、PRODUCT、MAX、MIN)
在 gather 和 reduce 概念前面还可以加上 all，如 all_gather，all_reduce，那就是多对多的关系了

在这里插入图片描述

多卡训练原理

网络在前向传播时会将 model 从主卡 (默认是逻辑 0 卡) broadcast 到所有 device 上，input data 会在 batch 这个维度被分组后 scatter 到不同的 device 上进行前向计算 (tuple, list and dict types will be shallow copied. The other types will be shared among different threads and can be corrupted if written to in the model’s forward pass.)，计算完毕后网络的输出被 gather 到主卡上，loss 随后在主卡上被计算出来 (这也是为什么主卡负载更大的原因，loss 每次都会在主卡上计算，这就造成主卡负载远大于其他显卡)。在反向传播时，loss 会被 scatter 到每个 device 上，每个卡各自进行反向传播计算梯度，然后梯度会被 reduce 到主卡上 (i.e. 求得各个 device 的梯度之和，然后按照 batch_size 大小求得梯度均值)，再用反向传播在主卡上更新模型参数，最后将更新后的模型参数 broadcast 到其余 GPU 中进行下一轮的前向传播，以此来实现并行

`nn.DataParallel()` 的用法

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

module (Module): 要放到多卡训练的模型
device_ids (list of python:int or torch.device): 可用的 gpu 卡号
output_device (int or torch.device): 模型输出结果存放的卡号 (如果不指定的话,默认放在 0 卡,这也是为什么多 gpu 训练并不是负载均衡的, 一般 0 卡负载更大)
dim (int)：从哪一维度切分一个 batch 的数据，默认为 0，即从 batch 维度将数据分组后送到不同 device 上运算

Example

nn.DataParallel() 的用法十分简单，加一行代码即可

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2]).cuda()	# broadcast model parameters to all devices
output = net(input_data)  # input_var can be on any device, including CPU

事实上 DataParallel 也是一个 Pytorch 的 nn.Module，因此使用 nn.DataParallel 后，模型需要使用 .module 来得到实际的模型和优化器：

# 保存模型
torch.save(net.module.state_dict(), path)

Use `nn.parallel.DistributedDataParallel` instead of `nn.DataParallel`

DataParallel 复制一个网络到多个 cuda 设备，然后再 split 一个 batch 的 data 到多个 cuda 设备，通过这种并行计算的方式解决了 batch 很大的问题，但也有自身的不足：
- (1) 单进程多线程带来的问题：DataParallel 是单进程多线程的，无法在多个机器上工作 (不支持分布式)，而且不能使用 Apex 进行混合精度训练。同时它基于多线程的方式，确实方便了信息的交换，但受困于 GIL (Python 全局解释器锁)，会带来性能开销 (GIL 的存在使得一个 Python 进程只能利用一个 CPU 核心，不适合用于计算密集型的任务)
- (2) 存在效率问题，主卡性能和通信开销容易成为瓶颈，GPU 利用率通常很低：数据集需要先拷贝到主进程，然后再 split 到每个设备上；权重参数只在主卡上更新，需要每次迭代前向所有设备做一次同步；每次迭代的网络输出需要 gather 到主卡上
- (3) 不支持 model parallel
这个时候，DistributedDataParallel 来了，并且自此之后，不管是单机还是多机，都推荐使用 DDP 来代替 DP. DP 和 DDP 的主要差异可以总结为以下几点：
- (1) DP 是单进程多线程的，只用于单机情况，而 DDP 是多进程的，每个 GPU 对应一个进程，适用于单机和多机情况，真正实现分布式训练，并且因为每个进程都是独立的 Python 解释器，DDP 避免了 GIL 带来的性能开销
- (2) DDP 的训练更高效，不存在 DP 中负载不均衡的问题，基本上 DP 已经被弃用；DDP 中每个 GPU 直接处理 mini-batch 数据，不需要由主卡分发；每个 GPU 独立进行参数更新，不需要由主卡 broadcast 模型参数；每个 GPU 独立计算 loss，不需要汇聚到主卡计算
- (3) DDP 支持模型并行，而 DP 并不支持，这意味如果模型太大单卡显存不足时只能使用前者

可能 DDP 唯一不好的地方就是相比 DP 使用起来会有些麻烦

`nn.parallel.DistributedDataParallel` (DDP)

分布式训练常见概念

node：节点，可以看作主机
global_rank：表示进程序号，用于进程间通信，可以用于表示进程的优先级。我们一般设置 global_rank=0 对应的主机为 master 节点
local_rank：主机内 GPU 编号，非显式参数，由 torch.distributed.run 内部指定 (torch.distributed.run 是为了代替 torch.distributed.launch 的新型启动方式，但是由于是新功能, 只有最新的 torch 1.10 支持, 因此出于兼容性考虑还是建议使用 torch.distributed.launch)。比方说，node_rank=3，local_rank=0 表示第 3 个主机内的第 1 块 GPU，因此 local_rank 对应的就是 Process 需要使用的 Device (GPU) 编号
world_size：全局进程个数 (在 DDP 中，一个进程控制一个 GPU，因此 world_size 即为 nnode $\times$ nproc_per_node (设备数)

DDP 内部机制

每个 GPU 都由一个进程控制，这些 GPU 可以都在同一个节点上 (单机)，也可以分布在多个节点上 (多机)。每个进程都执行相同的任务，并且每个进程都与所有其他进程通信。进程或者说 GPU 之间只传递梯度，这样网络通信就不再是瓶颈
首先将 rank=0 进程中的模型参数 (i.e. state_dict()) broadcast 到进程组中的其他进程 (当然，模型不止有参数，还有 buffer，它们的优化不是通过梯度反向传播而是通过其他方式，例如 BN 中的 moving mean and var)，然后每个 DDP 进程都会创建一个 local Reducer 来负责梯度同步。在训练过程中，每个进程从磁盘加载 batch 数据，并将它们传递到其 GPU。每个 GPU 都有自己的前向和反向过程，完成前向和反向传播后梯度在各个 GPUs 间进行 All-Reduce，每个 GPU 都收到其他 GPU 的梯度，从而可以独自进行参数更新。同时，每一层的梯度不依赖于前一层，所以梯度的 All-Reduce 和反向过程同时计算 (参数被分组为了多个 bucket，先计算、得到梯度的 bucket 会马上进行通讯，不必等到所有梯度计算结束才进行通讯)，以进一步缓解网络瓶颈 (DDP 把 Ring-Reduce 的代码写在了 Pytorch 中反向梯度计算结束后调用的 hook 接口函数里，这样每次 parameter 的反向梯度计算结束后，程序就会调用这个 hook 函数，从而开启 Ring-Reduce 流程)。在后向过程的最后，每个节点都得到了平均梯度，这样各个 GPU 中的模型参数保持同步 (回忆一下，DP 是将梯度 reduce 到主卡，在主卡上更新参数，再将参数 broadcast 给其他 GPU，这样无论是主卡的负载还是通信开销都比 DDP 大很多) (具体的实现细节可以参考 DISTRIBUTED DATA PARALLEL)
上述过程要求多个进程 (可能在多个节点上) 同步并通信。Pytorch 通过 distributed.init_process_group 函数初始化进程组来实现多进程同步通信。它需要知道 rank 0 位置以便所有进程都可以同步，以及预期的进程总数 (world_size)。每个进程都需要知道进程总数及其在进程组中的顺序，以及使用哪个 GPU. 此外，Pytorch 还提供了 torch.utils.data.DistributedSampler 来为各个进程切分数据，以保证训练数据不重叠

Ring All-reduce

在 GPU backend NCCL 中，all reduce 的实现方法就是 Ring All-reduce. 在 Ring All-reduce 中，所有设备组成一个环形，每个进程只跟自己上下游两个进程进行通讯，极大地缓解了参数服务器的通讯阻塞现象
具体来说，Ring All-reduce 由 Ring Reduce-scatter 和 Ring All-gather 两个步骤组成. 首先，需要把数据切成 $n$ 块，Ring Reduce-scatter 经过 $n - 1$ 步完成
Ring All-gather 经过 $n - 1$ 步完成

`nn.parallel.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, gradient_as_bucket_view=False, 
	static_graph=False)

module (Module): 要放到多卡训练的模型
device_ids (list of python:int or torch.device): 训练使用的 device id. 对于一卡一进程的情况，device_ids 设为 [local_rank]
output_device (int or torch.device): 模型输出结果存放的卡号。对于一卡一进程的情况，output_device 设为 local_rank
dim (int)：从哪一维度切分一个 batch 的数据，默认为 0，即从 batch 维度将数据分组后送到不同 device 上运算
其余参数保持默认即可，参数详情可参考 torch.nn.parallel.DistributedDataParallel

分布式接口

`dist.barrier()` - 同步所有进程

Synchronizes all processes. 调用该函数可以阻塞进程组内的所有进程，直至它们都进入了该函数

dist.barrier()

`all_reduce`

Reduces the tensor data across all machines in such a way that all get the final result. After the call tensor is going to be bitwise identical in all processes. The function operates in-place.

# define tensor on GPU, count and total is the result at each GPU
t = torch.tensor([count, total], dtype=torch.float64, device='cuda')

dist.all_reduce(t, op=dist.ReduceOp.SUM)

`all_gather`

Gathers tensors from the whole group in a list.

# define tensor on GPU, count and total is the result at each GPU
t = torch.tensor([count, total], dtype=torch.float64, device='cuda')

t_gather_list = [torch.zeros_like(t) for _ in range(world_size)]
dist.all_gather(t_gather_list, t)

`broadcast`

Broadcasts the tensor to the whole group.

# define tensor on GPU, count and total is the result at each GPU
t = torch.tensor([count, total], dtype=torch.float64, device='cuda')

# "t" is the data to be sent if src is the rank of current process, 
# and "t" to be used to save received data otherwise.
torch.distributed.broadcast(t, src)

`reduce`

Reduces the tensor data across all machines. Only the process with rank dst is going to receive the final result.

# define tensor on GPU, count and total is the result at each GPU
t = torch.tensor([count, total], dtype=torch.float64, device='cuda')

dist.reduce(t, dst, op=dist.ReduceOp.SUM)

`reduce_scatter`

Reduces, then scatters a list of tensors to all processes in a group.

# input_list 是 Tensor 的列表。先对 input_list 进行 all_reduce，得到的结果再 scatter 到各个进程中，
# 每个进程得到的输出结果即为 output
dist.reduce_scatter(output, input_list, op=dist.ReduceOp.SUM)

其他常用 API

# get world size
dist.get_world_size()

# get global rank
# global rank 为 0 的进程就是 master 进程
dist.get_rank()

如果使用 torch.distributed.launch / torchrun 启动分布式训练，则也可以通过环境变量获取分布式参数

os.environ["MASTER_ADDR"]
os.environ["MASTER_PORT"]
int(os.environ["RANK"])		# global rank
int(os.environ['LOCAL_RANK'])
int(os.environ["WORLD_SIZE"])

使用范例

导入需要的库

import os
import torch
import argparse
import torch.optim as optim
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
from datetime import timedelta

初始化进程组、配置 Master Node

init_process_group 参数详见 torch.distributed.init_process_group)

init_process_group. 各个进程会在这一步，与 master 节点进行握手，建立连接。如果连接上的进程数量不足约定的 word_size，进程会一直等待。也就是说，如果你约定了 world_size=64，但是只开了 6 台 8 卡机器，那么程序会一直暂停在这个地方
参数 init_method 包括 ‘env://’ (常用)、‘tcp://ip:port’ 和 ‘file://path’，用来表示 where/how to discover peers.
- ‘env://’ 表示从环境变量中读取分布式信息，主要包括 MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE，其中 rank 和 world_size 可以选择手动指定，否则从环境变量读取
- ‘tcp://ip:port’ 通过指定 rank 0 (即 MASTER 进程) 的 IP 和端口，各个进程进行信息交换。需指定 rank 和 world_size 这两个参数
- ‘file://path’ 通过所有进程都可以访问的共享文件系统来进行信息共享。需要指定 rank 和 world_size 参数

def setup():
	torch.distributed.init_process_group(
		backend="nccl",
		init_method='env://', 	# indicates where/how to discover peers
								# 'env://' means environment variable initialization
								# 'file:///mnt/nfs/sharedfile' means shared file-system initialization
								# 'tcp://10.1.1.20:23456' means TCP initialization
		timeout=timedelta(minutes=60)
	)

def cleanup():
    dist.destroy_process_group()

下面给出用 TCP 方式进程组初始化的示例

dist.init_process_group('nccl', init_method='tcp://127.0.0.1:28765', rank=args.rank, world_size=args.ws)

下面给出用 TCP 方式或 ENV 方式初始化进程组时，程序的运行方式 (对于 ENV 方式而言，torch.distributed.launch 在运行进程时会自动设置相应的环境变量，不需要手动去添加了)

# TCP 方法
python3 test_ddp.py --init_method=TCP --rank=0 --ws=2
python3 test_ddp.py --init_method=TCP --rank=1 --ws=2
# ENV 方法
MASTER_ADDR='localhost' MASTER_PORT=28765 RANK=0 WORLD_SIZE=2 python3 test_ddp.py --init_method=ENV
MASTER_ADDR='localhost' MASTER_PORT=28765 RANK=1 WORLD_SIZE=2 python3 test_ddp.py --init_method=ENV

定义训练过程

DDP 会在 model = DDP(model) 这一步把 parameter，buffer 从 master 节点传到其他节点，在每个进程上创建模型。需要注意的是，DDP 通过这一步保证所有进程的初始状态一致。所以，请确保在这一步之后，你的代码不会再修改模型的任何东西了，包括添加、修改、删除 parameter 和 buffer！
DistributedSampler. 在数据集采样时，给 dataloader 加一个 DistributedSampler 就可以给不同进程分配数据集的不重叠、不交叉的部分，从而无缝对接 DDP 模式 (其实就是 shuffle 数据集后，把样本依次分配给不同进程，例如样本 0 给进程 0，样本 1 给进程 1…)。那么问题来了，每次 epoch 我们都会随机 shuffle 数据集，那么，不同进程之间要怎么保持 shuffle 后数据集的一致性呢？DistributedSampler 的实现方式是，不同进程会使用一个相同的随机数种子，这样 shuffle 出来的东西就能确保一致。但如果在不同 epoch 也使用相同的随机数种子，那么每个 epoch 各个进程得到的数据集遍历顺序都是固定的，为了解决该问题，需要每次 epoch 开始前都要调用一下 sampler 的 set_epoch 方法，实际设置的 shuffle 随机数种子为 self.seed+self.epoch

def run_demo(args):
    # get global_rank and world_size, set device
    if args.local_rank == -1:
    	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
    	setup()
    	torch.cuda.set_device(args.local_rank)
    	device = torch.device("cuda", args.local_rank)

    # load model to the GPU specified by local_rank
    model = ToyModel().to(device)
    ddp_model = DDP(model, 
    				device_ids=[args.local_rank], 
    				output_device=args.local_rank)

	# define loss and optimizer
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

	# define data sampler and dataloader
	train_dataset = torchvision.datasets.MNIST(
	    root='./data',
	    train=True,
	    transform=transforms.ToTensor(),
	    download=True
	)
	train_sampler = torch.utils.data.distributed.DistributedSampler(
		train_dataset, 
		shuffle=True)
	trainloader = DataLoader(train_dataset,
							batch_size=args.batch_size,
							num_workers=args.workers,
							sampler=train_sampler, 
							pin_memory=True)

	for epoch in range(args.epoch):
		# set epoch
		train_sampler.set_epoch(epoch)
		for i, data in enumerate(trainloader):
			# get the inputs
			inputs, labels = data
			inputs = inputs.to(device)
			labels = labels.to(device)

			# zero the parameter gradients
		    optimizer.zero_grad()

			# forward + backward + optimize
		    outputs = ddp_model(inputs)
		    loss = loss_fn(outputs, labels)
		    loss.backward()
		    optimizer.step()
		
		dist.barrier()

    cleanup()

注意，如果用 torchrun 而非 torch.distributed.launch 启动分布式训练，则不应添加命令行参数 local_rank，而是应该从环境变量中获取 local_rank

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int, default=-1, 
    					help='local_rank for distributed training on gpus.' 
    						 'Automatically set by torch.distributed.launch')
    parser.add_argument('--epoch', type=int, help='#epochs to train')
    parser.add_argument('--batch_size', type=int, help='batch size')
    parser.add_argument('--workers', type=int, help='#workers for dataloader')
    args = parser.parse_args()

    run_demo(args=args)

启动分布式训练

`torch.distributed.launch`

nproc_per_node: 一个节点中进程的数量 (一般一张 GPU 对应一个进程)
nnodes: 节点的数量，通常一个节点对应一个主机，方便记忆，直接表述为主机，默认为 1
node_rank: 节点的序号，从 0 开始，默认为 0
master_addr: master 节点的 ip 地址。多机模式才会用到，默认为 127.0.0.1 (单机)
master_port: master 节点的 port 号，在不同的节点上 master_addr 和 master_port 的设置是一样的，用来进行通信。多机模式才会用到，需要保证该端口号没有被别的程序占用，默认为 29500

当使用多机分布式训练时，需要在每台机器上都运行一次 torch.distributed.launch

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch \
								   --nproc_per_node=2 \
								   --nnodes=1 \
								   --node_rank=0  \
								   --master_addr="10.103.10.54" \
								   --master_port=6005 \
								   train.py

如果是单机多卡的话，则只需要设置 nproc_per_node

torch.distributed.launch 在运行进程时会设置相应的环境变量 (对应 “env://” 的初始化方式)，在程序中可以通过环境变量获取相应设置

os.environ["MASTER_ADDR"]
os.environ["MASTER_PORT"]
int(os.environ["RANK"])	# global rank
int(os.environ["WORLD_SIZE"])

`torchrun`

torch1.10 开始用终端命令 torchrun 来代替 torch.distributed.launch，具体来说，torchrun 实现了 launch 的一个超集，不同的地方在于：

(1) 完全使用环境变量配置各类参数，如 RANK, LOCAL_RANK, WORLD_SIZE 等，尤其是 local_rank 不再支持用命令行隐式传递的方式.
- If your script expects --local_rank argument to be set, please change it to read from os.environ['LOCAL_RANK'] instead.
(2) 能够更加优雅地处理某个 worker 失败的情况，重启 worker。如果代码中有 load_checkpoint(path) 和 save_checkpoint(path) 这样有 worker 失败的话，可以通过 load 最新的模型，重启所有的 worker 接着训练
(3) 训练的节点数目可以弹性变化

torchrun 的使用

如果想要从 torch.distributed.launch 迁移到 torchrun，只需将程序中使用到的 local_rank 由从命令行参数获取改为从环境变量获取

local_rank = int(os.getenv('LOCAL_RANK', -1))

更多参数设置可参考 TORCHRUN (ELASTIC LAUNCH)

CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 \
								  train.py

`torch.multiprocessing.spawn`

…

保存模型参数

和 DP 一样，保存的是 model.module 而不是 model
同时注意只在 master 进程保存模型

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

其他注意事项

模型参数的定义顺序

在 DDP 中，参数被划分为了多个 bucket，当某个 bucket 的所有 parameter 都梯度都计算好时，reducer 就会开始对这个 bucket 的所有 parameter 进行异步的 all-reduce 梯度平均操作，而不必等到所有参数梯度全部计算完毕。但 bucket 的执行过程也是有顺序的，其顺序与 parameter 的定义是相反的，即最后注册的 parameter 的 bucket 在最前面，所以，我们在创建 module 的时候，请务必把先进行计算的 parameter 注册在前面，后计算的在后面 (这样后计算的参数，即梯度反向传播时最先计算出的参数对应的 bucket 就在最前面，能加快梯度同步速度)。不然，reducer 会卡在某一个 bucket 等待，使训练时间延长！

在 DDP 中引入 SyncBN

Why SyncBN?

SyncBN 就是 BN 的并行版本，SyncBN 能够完美支持多卡训练，而普通 BN 在多卡模式下实际上就是单卡模式
我们知道，BN 中有 moving mean 和 moving variance 这两个 buffer，这两个 buffer 的更新依赖于当前训练轮次的 batch 数据的计算结果。但是在普通多卡 DP 模式下，各个模型只能拿到自己的那部分计算结果，所以在 DP 模式下的普通 BN 被设计为只利用主卡上的计算结果来计算 moving mean 和 moving variance，之后再广播给其他卡。这样，实际上 BN 的 batch size 就只是主卡上的 batch size 那么大。当模型很大、batch size 很小时，这样的 BN 无疑会限制模型的性能
为了解决这个问题，PyTorch 新引入了一个叫 SyncBN 的结构，利用 DDP 的分布式计算接口来实现真正的多卡 BN

SyncBN 的原理

SyncBN 利用分布式通讯接口在各卡间进行通讯，从而能利用所有数据进行 BN 计算。为了尽可能地减少跨卡传输量，SyncBN 做了一个关键的优化，即只传输各自进程的各自的小 batch mean 和小 batch variance，而不是所有数据。具体流程如下：
(1) 在各进程上分别计算各自的小 batch mean 和小 batch variance；(2) 各自的进程对各自的小 batch mean 和小 batch variance 进行 all_gather 操作；(3) 每个进程由 all_gather 得到的数据分别计算总体 mean 和总体 variance，得到与正常 BN 一样的结果

在 PyTorch 中使用 SyncBN

当前 PyTorch SyncBN 只在 DDP 单进程单卡模式中支持 (SyncBN 依赖于 all_gather，而这个分布式接口当前是不支持单进程多卡或者 DP 模式的。当然不排除未来也是有可能支持的)
利用 torch.nn.SyncBatchNorm.convert_sync_batchnorm 可以将模型中所有是或者继承了 torch.nn.modules.batchnorm._BatchNorm 类的普通 BN 都替换为 SyncBN，它需要加在模型定义后，创建 DDP 模型前 (因为创建 DDP 模型后，我们就不应该修改模型了)

dist.init_process_group(backend='nccl')

model = MyModel()
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

DDP 下 Gradient Accumulation 的进一步加速

为什么还能进一步加速？

我们仔细思考一下 DDP 下的 gradient accumulation. DDP 的 gradient all_reduce 阶段发生在 loss.backward()。这意味着，在梯度累加的情况下，假设一次梯度累加循环有 $K$ 个 step，每次梯度累加循环会进行 $K$ 次 all_reduce！但事实上，每次梯度累加循环只会有一次 optimizer.step()，即只应用一次参数修改，这意味着在每一次梯度累加循环中，我们其实只要进行一次 gradient all_reduce 即可满足要求，有 $K - 1$ 次 all_reduce 被浪费了！而每次 all_reduce 的时间成本是很高的！

for step, batch in enumerate(data_loader):
	loss = model(batch)
	loss = loss / K
	loss.backward()

	if (step + 1) % K == 0:
		optimizer.step()
    	optimizer.zero_grad()

如何加速

解决问题的思路在于，对前 $K - 1$ 次 step 取消其梯度同步。幸运的是，DDP 给我们提供了一个暂时取消梯度同步的 context 函数 no_sync()，在这个 context 下，DDP 不会进行梯度同步

# 兼容 DDP 和单卡模式 (local_rank == -1 表示单卡模式)
from contextlib import nullcontext

for step, batch in enumerate(data_loader):
	my_context = model.no_sync if local_rank != -1 and (step + 1) % K != 0 else nullcontext
	with my_context():
		loss = model(batch)
		loss = loss / K
		loss.backward()

	if (step + 1) % K == 0:
		optimizer.step()
    	optimizer.zero_grad()

多机多卡环境下的 inference 加速

利用 DDP 进行推理：各个进程中各自进行单卡的 inference，然后把结果收集到一起。问题其实只有两个：(1) 如何把数据 split 到各个进程中；(2) 如何把结果合并到一起
注意，本节的内容因应用场景而异，只是提供一种可能的解决方案

如何把数据 split 到各个进程中

在训练的时候，我们用的 torch.utils.data.distributed.DistributedSampler 帮助我们把数据不重复地分到各个进程上去。但是，其分的方法是：每段连续的 $N$ 个数据，拆成一个一个，分给 $N$ 个进程，所以每个进程拿到的数据不是连续的，这样不利于我们在 inference 结束的时候将结果合并到一起 (这点应该视具体应用场景而定，不一定需要每个进程分到连续数据块)
下面实现了一个新的 data sample，它能够连续地划分数据块，不重复地分到各个进程上

# from: https://github.com/huggingface/transformers/blob/447808c85f0e6d6b0aeeb07214942bf1e578f9d2/src/transformers/trainer_pt_utils.py
class SequentialDistributedSampler(torch.utils.data.sampler.Sampler):
    """
    Distributed Sampler that subsamples indicies sequentially,
    making it easier to collate all results at the end.
    Even though we only use this sampler for eval and predict (no training),
    which means that the model params won't have to be synced (i.e. will not hang
    for synchronization even if varied number of forward passes), we still add extra
    samples to the sampler to make it evenly divisible (like in `DistributedSampler`)
    to make it easy to `gather` or `reduce` resulting tensors at the end of the loop.
    """

    def __init__(self, dataset, batch_size, rank=None, num_replicas=None):
        if num_replicas is None:
            if not torch.distributed.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = torch.distributed.get_world_size()
        if rank is None:
            if not torch.distributed.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = torch.distributed.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.batch_size = batch_size
        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.batch_size / self.num_replicas)) * self.batch_size
        self.total_size = self.num_samples * self.num_replicas

    def __iter__(self):
        indices = list(range(len(self.dataset)))
        # add extra samples to make it evenly divisible
        indices += [indices[-1]] * (self.total_size - len(indices))
        # subsample
        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
        return iter(indices)

    def __len__(self):
        return self.num_samples

test_sampler = SequentialDistributedSampler(testset, batch_size=batch_size)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, sampler=test_sampler)

如何把结果合并到一起

def distributed_concat(tensor, num_total_examples):
    output_tensors = [tensor.clone() for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(output_tensors, tensor)
    concat = torch.cat(output_tensors, dim=0)
    # truncate the dummy elements added by SequentialDistributedSampler
    return concat[:num_total_examples]

for epoch in range(n_epoch):
    ...
    with torch.no_grad():
        # 1. 得到本进程的 prediction
        predictions = []
        labels = []
        for data, label in testloader:
            data, label = data.to(local_rank), label.to(local_rank)
            predictions.append(model(data))
            labels.append(label)
        # 2. 进行 gather
        predictions = distributed_concat(torch.concat(predictions, dim=0), 
                                         len(test_sampler.dataset))
        labels = distributed_concat(torch.concat(labels, dim=0), 
                                    len(test_sampler.dataset))
        # 3. 现在我们已经拿到所有数据的 predictioin 结果，进行 evaluate！
        my_evaluate_func(predictions, labels)

保证 DDP 性能：确保数据的一致性

性能期望：从原理上讲，当没有开启 SyncBN 时，(或者更严格地讲，没有 BN 层；但一般有的话影响也不大)，以下两种方法训练出来的模型应该是性能相似的：(1) 进程数为 $N$ 的 DDP 训练；(2) gradient accumulation 为 $N$ 、其他配置完全相同的单卡训练
如果我们发现性能对不上，那么往往是 DDP 中的某些设置出了问题。其中最有可能的是数据方面出现了问题。DDP 训练时，数据的一致性必须被保证：各个进程拿到的数据，要像是 accumulation 为 $N$ 、其他配置完全相同的单卡训练中同个 accumulation 循环中不同 iteration 拿到的数据。想象一下，如果各个进程拿到的数据是一样的，或者分布上有任何相似的地方，那么，这就会造成训练数据质量的下降，最终导致模型性能下降

保证数据相似性：随机数种子

为保证实验的可复现性，一般我们会在代码开头声明一个固定的随机数种子。但是在 DDP 训练中，如果还是像以前一样在代码中用固定值作为随机数种子，就会造成以下后果：
- (1) DDP 的 $N$ 个进程都使用同一个随机数种子
- (2) 在生成数据时，如果我们使用了一些随机过程的数据扩充方法，那么，各个进程生成的数据会带有一定的同态性。比如说，YOLOv5 会使用 mosaic 数据增强 (从数据集中随机采样 3 张图像与当前的拼在一起，组成一张里面有 4 张小图的大图)。这样，因为各卡使用了相同的随机数种子，你会发现，各卡生成的图像中，除了原本的那张小图，其他三张小图都是一模一样的！
- (3) 同态性的数据，降低了训练数据的质量，也就降低了训练效率！最终得到的模型性能，很有可能是比原来更低的
所以，我们需要给不同的进程分配不同的、固定的随机数种子

set_random_seed(seed=42 + local_rank)

控制不同进程的执行顺序

一般情况下，各个进程是各自执行的，速度有快有慢，只有在 gradient all-reduce 的时候才会进行进程同步。那么，如果我们需要在其他地方进行同步呢？比如说，在加载数据前，如果数据集不存在，我们要下载数据集：我们只需要在唯一一个进程中开启一次下载，然后让其他进程等待其下载完成，再去加载数据。为此，可以使用 dist.barrier()

只在主进程执行，无须同步

if local_rank == 0:
    code_only_run_in_main_process()

简单的同步

code_before()
dist.barrier()
code_after()

在某个进程中执行 A 操作，其他进程等待其执行完成后再执行 B 操作

if local_rank == 0:
    do_A()
    dist.barrier()
else:
    dist.barrier()
    do_B()

在某个进程中优先执行 A 操作，其他进程等待其执行完成后再执行 A 操作

这个值得深入讲一下，因为这个是非常普遍的需求。利用 contextlib.contextmanager，我们可以把这个逻辑给优雅地包装起来！

关于 yield 和 @contextmanager 可以参考 python 中 yield 的用法详解——最简单，最清晰的解释和 Python 标准模块 – ContextManager

from contextlib import contextmanager

@contextmanager
def torch_distributed_zero_first(rank: int):
    """Decorator to make all processes in distributed training wait for each local_master to do something.
    """
    if rank not in [-1, 0]:
        torch.distributed.barrier()
    yield
    if rank == 0:
        torch.distributed.barrier()

with torch_distributed_zero_first(rank):
    if not check_if_dataset_exist():
        download_dataset()
    load_dataset()

避免 DDP 带来的冗余输出

问题：当我们在自己的模型中加入 DDP 模型时，第一的直观感受肯定是，终端里的输出变成了 $N$ 倍了。这是因为我们现在有 $N$ 个进程在同时跑整个程序
解法：logging 模块 + 输出信息等级控制。即用 logging 输出代替所有 print 输出，并给不同进程设置不同的输出等级，只在 0 号进程保留低等级输出。举一个例子：

import logging

# 给主要进程（rank=0）设置低输出等级，给其他进程设置高输出等级。
logging.basicConfig(level=logging.INFO if rank in [-1, 0] else logging.WARN)
# 普通 log，只会打印一次。
logging.info("This is an ordinary log.")
# 危险的 warning、error，无论在哪个进程，都会被打印出来，从而方便 debug
logging.error("This is a fatal log!")

其他问题

关于 batch_size 参数和 real batch size

batch_size 参数是 per-process concept
由于 DP 是单进程多线程，因此如果设置 batch_size 参数为 $N$ ，则 real batch size 即为 $N$ ，每个 GPU 负责处理的 batch size 为 $num_of_devices N/\text{num\_of\_devices}$ (主卡负责把 mini-batch 分发给每个设备)
而由于 DDP 是多进程，因此如果设置 batch_size 参数为 $N$ ，则 real batch size 即为 $world_size N\times\text{world\_size}$ ，每个 GPU 负责处理的 batch size 为 $N$ (每个设备直接处理 mini-batch). 不过由于各个设备上 all reduce 得到的梯度会除以 world_size (average (instead of sum) gradients across processes)，因此学习率等参数都可以按照 batch_size 来调 ( $N$ 卡的 DDP 模式，理论上可以等价于 $N$ 次 gradient accumulation 的单卡模式)