Pytorch DistributedDataParallel 多卡训练

Pytorch DistributedDataParallel 多卡训练

这里首先列举一下借鉴的博客：

1. https://zhuanlan.zhihu.com/p/178402798 有三篇，由浅入深
2. https://zhuanlan.zhihu.com/p/76638962 介绍了很多常用函数和对应参数的功能

DistributedDataParallel 多卡可以分为数据并行（常用方式，每张卡具有相同的模型和参数，训练时将 batch 数据拆分输入不同的模型中），模型并行（将模型拆分，不同部分放置在不同的 GPU 上，并行计算），Workload Partitioning（将模型拆分，不同部分放置在不同的 GPU 上，串行计算）。

由于第一种数据并行方式是最常用，并且是官方实践性能最佳的方式，因此本文相关内容主要针对数据并行方式给出。

基础知识

GPU 并行计算基础处理流程：设置可见 GPU 并配置并行化参数，创建模型，通过 API 将模型并行化，将模型和数据搬到 GPU，进行前向和后向传播。在前向传播中，会自动将 batch_size 切分后分配到可见的 GPU 上并行计算。结束后，会有一台或其他方式收集前向传播计算结果并根据 loss 更新每块 GPU 上模型参数。

因此在多 GPU 训练时，可以提高整体 batch_size，同时增加 learning_rate（一般为 batch_size 增加倍数的一半）。

管理方式

• group：进程组，默认情况下，只有一个组，一个 job 即为一个组，也即一个 world。
• world size：表示全局进程个数，一般和 GPU 数相同（单进程单GPU情况）。
• rank：表示进程序号，用于进程间通讯，表征进程优先级，序号一般从 0 到 world_size - 1。rank = 0 的主机为 master 节点。
• local_rank：进程内 GPU 编号，非显式参数，一般为一台主机内的 GPU 序号（从 0 到该机 GPU 数减一），由 torch.distributed.launch 内部指定。

常用方法

torch.cuda

• torch.cuda.is_available()：判断 GPU 是否可用
• torch.cuda.device_count()：计算当前可见可用的 GPU 数
• torch.cuda.get_device_name()：获取 GPU 型号，如 Tesla K80
• torch.cuda.manual_seed()：为当前 GPU 设置随机种子
• torch.cuda.manual_seed_all()：为所有可见可用 GPU 设置随机种子
• torch.cuda.current_device()：返回当前设备索引
• os.environ.setdefault("CUDA_VISIBLE_DEVICES", "2,3")：设置实际可见的 GPU，等价于 os.environ['CUDA_VISIBLE_DEVICES'] = '2,3'
• torch.cuda.set_device()：作用与 os.environ.setdefault() 类似，官方建议使用 os.environ.setdefault()。但实际这个更好使 =_=

torch.distributed

• torch.distributed.get_world_size()：获取全局并行数
• torch.distributed.new_group()：使用 world 的子集，创建新组，用于集体通信等
• torch.distributed.get_rank()：获取当前进程的序号，用于进程间通讯。
• torch.distributed.local_rank()：获取本台机器上的进程的序号

其他

• torch.device()：创建 device 对象，如 torch.device('cpu')，torch.device('cuda:1')
• tensor.to()，module.to()：将 tensor 转换类型或者搬到 GPU（会重新创建一个新的 tensor），将 module 搬到 GPU （会复用之前的 module）

torch.nn.parallel.DistributedDataParallel 介绍

首先给出几种并行化方案的性能对比图。

测试结果发现 Apex 的加速效果最好，但与 Horovod/Distributed 差别不大，平时可以直接使用内置的 Distributed。

而 torch.nn.DataParallel 效果不好的原因主要是其全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。而 DistributedDataParallel 在每次迭代中，每个进程具有自己的 optimizer ，并独立完成所有的优化步骤。在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。相较于 DataParallel，torch.distributed 传输的数据量更少，因此速度更快，效率更高。

同时 pytorch 官方文档也建议使用 DistributedDataParallel（DDP）替换 DataParallel（DP），因此本文着重介绍 DistributedDataParallel（DDP）的相关内容。

DDP 使用方式

DDP 多卡运行的本质还是通过创建多个进程来实现并行，但由于多个进程使用的是同一份代码，因此需要在代码中增加相关逻辑来指定进程与 GPU 硬件之间的关联关系。同时为了实现数据的并行化，需要为不同的进程（不同的 GPU）加载不同的数据，因此需要一个特殊的 data sampler 来实现，这个 DDP 通过 torch.utils.data.distributed.DistributedSampler 来实现。

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

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

模型并行化操作

DDP 启动方式存在 TCP 和环境变量启动两种方式，下面给出环境变量启动方式。

## 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) 实现了把parameter，buffer从master节点传到其他节点，使所有进程上的状态一致。所以，请确保在这一步之后，你的代码不会再修改模型的任何东西了，包括添加、修改、删除parameter和buffer！

其次 model 的定义时，请要求按照网络计算顺序，依次创建网络层。因为为了加速，DDP reducer 内部会根据注册的顺序反向判断各 GPU 中对应模型梯度是否已经全部求出。如果已经全部求出，那么就可以提前进行异步的 all-reduce 梯度平均操作。否则，reducer 会卡在某一个 bucket（将某些相邻的 parameter 组成 bucket 简化计算）等待，使训练时间延长！

因为 optimizer 和 DDP 是没有关系的，所以 optimizer 初始状态的同一性是不被 DDP 保证的！不过，大多数官方optimizer，其实现能保证从同样状态的 model 初始化时，其初始状态是相同的。所以这边我们只要保证在 DDP 模型创建后才初始化 optimizer，就不用做额外的操作。但是，如果自定义optimizer，则需要你自己来保证其统一性！

数据并行化操作

my_trainset = torchvision.datasets.CIFAR10(root='./data', train=True)
# 新增1：使用DistributedSampler，DDP帮我们把细节都封装起来了。用，就完事儿！
#       sampler的原理，后面也会介绍。
train_sampler = torch.utils.data.distributed.DistributedSampler(my_trainset)
# 需要注意的是，这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
trainloader = torch.utils.data.DataLoader(my_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()

DistributedSampler的实现方式是，不同进程会使用一个相同的随机数种子，这样shuffle出来的东西就能确保一致。具体实现上，DistributedSampler使用当前epoch作为随机数种子，从而使得不同epoch下有不同的shuffle结果。所以，记得每次 epoch 开始前都要调 用一下 sampler 的 set_epoch 方法，这样才能让数据集随机 shuffle 起来。

模型参数保存

# 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")

启动方式

针对 DDP 并行化，需要用 torch.distributed.launch 来启动训练。其中可选参数如下

• --nnodes：可用的机器数
• --node_rank：当前机器的编号
• --nproc_per_node：每台机器启动的进程数，一般为 GPU 数
• address，port：多机多卡是通过这个指定的 IP 和端口号进行数据通信

## Bash运行
# 假设我们只在一台机器上运行，可用卡数是8
python -m torch.distributed.launch --nproc_per_node 8 main.py


# 假设我们在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

这里需要注意：如果连接上的进程数量不足约定的 word_size，进程会一直等待。同时 torch.multiprocessing.spawn 可以对 DDP 进行封装从而实现基于单卡相同的方式启动。

分布式相关 API 介绍

init_process_group

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

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

• backend ：指定当前进程要使用的通信后端。小写字符串，支持的通信后端有 gloo，mpi，nccl 。建议用 nccl。
• init_method ： 指定当前进程组初始化方式。可选参数，字符串形式。如果未指定 init_method 及 store，则默认为 env://，表示使用读取环境变量的方式进行初始化。该参数与 store 互斥。
• rank： 指定当前进程的优先级。int 值。表示当前进程的编号，即优先级。如果指定 store 参数，则必须指定该参数。rank=0 的为主进程，即 master 节点。
• world_size：该 job 中的总进程数。如果指定 store 参数，则需要指定该参数。
• timeout ： 指定每个进程的超时时间。可选参数，datetime.timedelta 对象，默认为 30 分钟。该参数仅用于 Gloo 后端。
• store：所有 worker 可访问的 key / value，用于交换连接 / 地址信息。与 init_method 互斥。

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 对象，用于指定要并行的设备。对于数据并行，即完整模型放置于一个 GPU 上（single-device module）时，需要提供该参数，表示将模型副本拷贝到哪些 GPU 上。对于模型并行的情况，即一个模型，分散于多个 GPU 上的情况（multi-device module），以及 CPU 模型，该参数必须为 None，或者为空列表。与单机并行一样，输入数据及中间数据，必须放置于对应的，正确的 GPU 上。
• output_device：int 或者 torch.device。对于 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)

DistributedSampler

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

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

注意事项

1、多 GPU 训练模型保存时会将模型参数保存在 module 域下，如 module.linears.0.weight

torch.save() 和 torch.load() 函数本质将模型参数转存为一个 OrderDict 字典。但当模型是在多 GPU 方式下训练时，标识对应模型参数的 key 会自动添加上 module 中间层。这在重新加载模型时可能造成错误，可以使用如下代码去除 module 层

from collections import OrderedDict

new_state_dict = OrderedDict()
for k, v in state_dict_load.items():
    namekey = k[7:] if k.startswith('module.') else k
    new_state_dict[namekey] = v

2、想用更大的 batch_size 进行训练，可以使用梯度累加

梯度累加的基本思想在于，在优化器更新参数前，也就是执行 optimizer.step() 前，进行多次反向传播，是的梯度累计值自动保存在 parameter.grad 中，最后使用累加的梯度进行参数更新。这个在 PyTorch 中特别容易实现，因为 PyTorch 中，梯度值本身会保留，除非我们调用 model.zero_grad() 或 optimizer.zero_grad()。

model.zero_grad()                                   # 重置保存梯度值的张量

for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # 前向计算
    loss = loss_function(predictions, labels)       # 计算损失函数
    loss.backward()                                 # 计算梯度
    if (i + 1) % accumulation_steps == 0:           # 重复多次前面的过程
        optimizer.step()                            # 更新梯度
        model.zero_grad()                           # 重置梯度

3、DDP 中设置的 batch_size 就是这个进程 forward 时使用的 batch_size，因此应该根据 world_size 和 整体需要的 batch_size 来调整每个 DDP 中设置的 batch_size

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

world_size = dist.get_world_size()
assert total_batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count'

train_loader = torch.utils.data.DataLoader(dataset=data_train,
                                                    batch_size=total_batch_size // world_size,
                                                    # shuffle=True,
                                                    num_workers=2,
                                                    sampler=train_sampler)

4、由于 BN 层需要基于传入模型的数据计算均值和方差，造成普通 BN 在多卡模式下实际上就是单卡模式。此时需要使用 SyncBN 利用DDP的分布式计算接口来实现真正的多卡BN。

SyncBN利用分布式通讯接口在各卡间进行通讯，传输各自进程小 batch mean 和小 batch variance，在传输少量数据的基础上利用所有数据进行BN计算。不过，当前PyTorch SyncBN只在DDP单进程单卡模式中支持。

同时由于 SyncBN 用到 all_gather 这个分布式计算接口，而使用这个接口需要先初始化DDP环境，因此 SyncBN 需要在 DDP 环境初始化后初始化，但是要在 DDP 模型前就准备好。

最后由于 SyncBN 是直接搜索 model 中每个 module，如果这个 module 是 torch.nn.modules.batchnorm._BatchNorm 的子类，就将其替换为 SyncBN。因此如果你的 Normalization 层是自己定义的特殊类，没有继承过 _BatchNorm 类，那么convert_sync_batchnorm 是不支持的，需要你自己实现一个新的SyncBN！

# DDP init
dist.init_process_group(backend='nccl')

# 按照原来的方式定义模型，这里的BN都使用普通BN就行了。
model = MyModel()
# 引入SyncBN，这句代码，会将普通BN替换成SyncBN。
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

# 构造DDP模型
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

5、使用 gradient accumulation 时，减少每一次梯度累加时的无效 all_reduce 从而实现 DDP 的进一步加速。

from contextlib import nullcontext
# 如果你的python版本小于3.7，请注释掉上面一行，使用下面这个：
# from contextlib import suppress as nullcontext

if local_rank != -1:
    model = DDP(model)

optimizer.zero_grad()
for i, (data, label) in enumerate(dataloader):
    # 只在DDP模式下，轮数不是K整数倍的时候使用no_sync
    my_context = model.no_sync if local_rank != -1 and i % K != 0 else nullcontext
    with my_context():
        prediction = model(data)
        loss_fn(prediction, label).backward()
    if i % K == 0:
        optimizer.step()
        optimizer.zero_grad()

DDP 的梯度 all_reduce 发生在 loss.backward() 时，而由于在使用梯度累加时，K 次 backward 之后才会真正使用更新后的梯度，因此前 K-1 次 all_reduce 其实都是无效的。因此可以通过 DDP 给我们提供的 no_sync 函数暂时取消梯度同步，从而进一步加速性能。

6、随机数种子

import random
import numpy as np
import torch

def init_seeds(seed=0, cuda_deterministic=True):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
    if cuda_deterministic:  # slower, more reproducible
        cudnn.deterministic = True
        cudnn.benchmark = False
    else:  # faster, less reproducible
        cudnn.deterministic = False
        cudnn.benchmark = True
        

def main():
    rank = torch.distributed.get_rank()
    # 问题完美解决！
    init_seeds(1 + rank)

在实验过程中应该避免 DDP 的每一个进程使用相同的随机数种子，因为这会造成各个进程生成的数据带有一定的同态性，从而降低训练数据的质量和训练效率。

这个和数据并行化时，使用 trainloader.sampler.set_epoch(epoch) 道理一致，这里是为了保证不同的 epoch 间不同进程将得到不同的训练数据。

7、日志的简化输出

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!")

这个比较直接，通过配置不同的日志输出等级，从而避免每个进程频繁的输出日志。