PyTorch 单机多GPU 训练方法与原理整理

这里整理一些PyTorch单机多核训练的方法和简单原理，目的是既能在写代码时知道怎么用，又能从原理上知道大致是怎么回事儿。

就目前来说，并行训练的方法可以根据的不同的并行对象分为——模型并行和数据并行。

模型并行：是网络太大，一张卡存不了，那么拆分，然后进行模型并行训练。
数据并行：多个显卡同时采用数据训练网络的副本。
在这里仅先讨论数据并行。

PyTorch单机多核训练方案有两种：一种是利用nn.DataParallel实现，实现简单，不涉及多进程；另一种是用torch.nn.parallel.DistributedDataParallel和torch.utils.data.distributed.DistributedSampler结合多进程实现。第二种方式效率更高，但是实现起来稍难，第二种方式同时支持多节点分布式实现。方案二的效率要比方案一高，即使是在单运算节点上。

并行数据加载

先介绍一下数据并行的概念。

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

将数据从磁盘加载到主机；
将数据从可分页内存传输到主机上的固定内存；
将数据从固定内存传输到GPU；
在GPU上向前和向后传递；

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

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

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

接下来将介绍具体的实现方法，为了方便理解，这里用一个简单的CNN模型训练MNIST手写数据集，相关代码：

model.py：定义一个简单的CNN网络
data.py：MNIST训练集和数据集准备
single_gpu_train.py：单GPU训练代码

方案一

核心在于使用nn.DataParallel将模型wrap一下，代码其他地方不需要做任何更改:

model = nn.DataParallel(model)

为方便说明，我们假设模型输入为(32, input_dim)，这里的 32 表示batch_size，模型输出为(32, output_dim)，使用 4 个GPU训练。nn.DataParallel起到的作用是将这 32 个样本拆成 4 份，发送给 4 个GPU 分别做 forward，然后生成 4 个大小为(8, output_dim)的输出，然后再将这 4 个输出都收集到cuda:0上并合并成(32, output_dim)。

详细流程：
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，所以他的内存和使用率会比其他的高。
具体流程见下图：

可以看出，nn.DataParallel没有改变模型的输入输出，因此其他部分的代码不需要做任何更改，非常方便。但弊端是，后续的loss计算只会在cuda:0上进行，没法并行，因此会导致负载不均衡的问题。

如果把loss放在模型里计算的话，则可以缓解上述负载不均衡的问题，示意代码如下：

class Net:
    def __init__(self,...):
        # code
    
    def forward(self, inputs, labels=None)
        # outputs = fct(inputs)
        # loss_fct = ...
        if labels is not None:
            loss = loss_fct(outputs, labels)  # 在训练模型时直接将labels传入模型，在forward过程中计算loss
            return loss
        else:
            return outputs

按照我们上面提到的模型并行逻辑，在每个GPU上会计算出一个loss，这些loss会被收集到cuda:0上并合并成长度为 4 的张量。这个时候在做backward的之前，必须对将这个loss张量合并成一个标量，一般直接取mean就可以。这在Pytorch官方文档 nn.DataParallel函数中有提到：

When module returns a scalar (i.e., 0-dimensional tensor) in forward(), this wrapper will return a vector of length equal to number of devices used in data parallelism, containing the result from each device.

这部分的例子可以参考：data_parallel_train.py

Note: DataParallel 中，batch size 设置必须为单卡的 n 倍，因为一个batch的数据会被主GPU分散为minibatch给其他GPU，但是在 DistributedDataParalle(方案二) 内，batch size 设置于单卡一样即可，因为各个GPU对应的进程独立从磁盘中加载数据，训练所设置的batch size就是每一个进程的batch size。

从一开始就会启动多个进程(进程数等于GPU数)，每个进程独享一个GPU，每个进程都会独立地执行代码。这意味着每个进程都独立地初始化模型、训练，当然，在每次迭代过程中会通过进程间通信共享梯度，整合梯度，然后独立地更新参数。
每个进程都会初始化一份训练数据集，当然它们会使用数据集中的不同记录做训练，这相当于同样的模型喂进去不同的数据做训练，也就是所谓的数据并行。这是通过torch.utils.data.distributed.DistributedSampler函数实现的，不过逻辑上也不难想到，只要做一下数据partition，不同进程拿到不同的parition就可以了，官方有一个简单的demo，感兴趣的可以看一下代码实现：Distributed Training
进程通过local_rank变量来标识自己，local_rank为0的为master，其他是slave。这个变量是torch.distributed包帮我们创建的，使用方法如下：
```
import argparse  # 必须引入 argparse 包
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()
```
必须以如下方式运行代码：
```
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 train.py
```
这样的话，torch.distributed.launch就以命令行参数的方式将args.local_rank变量注入到每个进程中，每个进程得到的变量值都不相同。比如使用 4 个GPU的话，则 4 个进程获得的args.local_rank值分别为0、1、2、3。

上述命令行参数nproc_per_node表示每个节点需要创建多少个进程(使用几个GPU就创建几个)；nnodes表示使用几个节点，因为我们是做单机多核训练，所以设为1。
Note:
1. 最新的pytorch提示，改为了torchrun，以单机多卡为例，启动方式为(更简单)：
```
torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py
# --nnodes表示为物理节点数量
# --nproc_per_node表示为GPU数量
```
1. torch.distributed.launch还有两个参数比较容易弄混，分别是**rank和local_rank**.
  1. rank：表现当前进程的序号，用于进程间通讯，表征进程优先级。rank = 0 的主机为 master 节点。同时，rank=0的进程就是master进程，其可以通过如下的方式获取:
```
# 获取rank，每个进程都有自己的序号，各不相同
torch.distributed.get_rank()
```
  1. local_rank：这是每台机子上的进程的序号，进程内，GPU 编号，非显式参数，由 torch.distributed.launch 内部指定。比方说， rank = 3，local_rank = 0 表示第 3 个进程内的第 1 块 GPU。获取方式如下：
```
# 获取local_rank。一般情况下，你需要用这个local_rank来手动设置当前模型
#是跑在当前机器的哪块GPU上面的。
torch.distributed.local_rank()
```
  2. 举个栗子：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)
2. 与 nn.DataParallel分发数据的方式不同，parallel.DistributedDataParallel的batch_size是不进行整除的，也即是你设置了BS是多少，每个进程(GPU)的BS就是多少，但是其会对数据进行划分(第一个维度)。
  
  举个例子来说：我们假设模型输入为(32, input_dim)，数据集大小为12000，有4块GPU，那么每块GPU上的BS大小为32，数据集大小为12000/32=375/4=93.74，近似为94。12000/32是说明，按照BS等于32来划分，总共有375个input，又因为有4块GPU，因此每块GPU的input的数量为94.
因为每个进程都会初始化一份模型，为保证模型初始化过程中生成的随机权重相同，需要设置随机种子。方法如下：
```
def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
```

使用方法通过如下示意代码展示：

from torch.utils.data.distributed import DistributedSampler  # 负责分布式dataloader创建，也就是实现上面提到的partition。

# 负责创建 args.local_rank 变量，并接受 torch.distributed.launch 注入的值
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()

# 每个进程根据自己的local_rank设置应该使用的GPU
torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)

# 初始化分布式环境，主要用来帮助进程间通信
torch.distributed.init_process_group(backend='nccl')

# 固定随机种子
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# 初始化模型
model = Net()
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 只 master 进程做 logging，否则输出会很乱
if args.local_rank == 0:
    tb_writer = SummaryWriter(comment='ddp-training')

# 分布式数据集
train_sampler = DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size)  # 注意这里的batch_size是每个GPU上batch_size

# 分布式模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)

详细代码参考：ddp_train.py

In distributed mode, calling the set_epoch() method at the beginning of each epoch before creating the DataLoader iterator is necessary to make shuffling work properly across multiple epochs. Otherwise, the same ordering will be always used.

使用方法很简单，只需要在input之间添加该方法就可以了:

sampler = DistributedSampler(dataset) if is_distributed else None
loader = DataLoader(dataset, shuffle=(sampler is None),
                    sampler=sampler)
for epoch in range(start_epoch, n_epochs):
    if is_distributed:
        sampler.set_epoch(epoch)
    train(loader)

torch.distributed.barrier

在读 huggingface/transformers中的源码，比如examples/run_ner.py会看到一下代码：

    # Load pretrained model and tokenizer
    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab
        # 这里有一点值得注意，不同进程的 barrier  实际上是互相对应的，必须所有进程都执行一次barrier，才会重新放行正常前进。

    args.model_type = args.model_type.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
                                          num_labels=num_labels,
                                          cache_dir=args.cache_dir if args.cache_dir else None)
    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
                                                do_lower_case=args.do_lower_case,
                                                cache_dir=args.cache_dir if args.cache_dir else None)
    model = model_class.from_pretrained(args.model_name_or_path,
                                        from_tf=bool(".ckpt" in args.model_name_or_path),
                                        config=config,
                                        cache_dir=args.cache_dir if args.cache_dir else None)

    if args.local_rank == 0:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

上述代码要实现预训练模型的下载和读入内存，如果4个进程都分别下载一遍显然是不合理的，那如何才能实现只让一个进程下载呢？这个时候就可以使用barrier函数。当slave进程(local_rank!=0)运行到第一个if时就被barrier住了，只能等着，但master进程可以往下运行完成模型的下载和读入内存，但在第二个if语句时遇到barrier，那会不会被barrier住呢？答案是不会，因为master进程和slave进程集合在一起了(barrier)，barrier会被解除，这样大家都往下执行。当然这时大家执行的进度不同，master进程已经执行过模型读入，所以从第二个if往下执行，而slave进程尚未执行模型读入，只会从第一个if往下执行。

可以看到barrier类似一个路障，进程会被拦住，直到所有进程都集合齐了才放行。适合这样的场景：只一个进程下载，其他进程可以使用下载好的文件；只一个进程预处理数据，其他进程使用预处理且cache好的数据等。

这里我遇到了一个大问题，就是比如我们想要做单一设备上的预测，不通过分布式的方法来预测，那么通过torch.distributed.barrier方法时，就会出现所有进程卡住。虽然此时的GPU占用率是100%，但是没有任何程序在正常运行，不过在我找了好多博客之后，终于在Github上发现了类似的问题Using torch.distributed.barrier() makes the whole code hang。
这里的方法简单来说，就是将原来的model(**input)用model.module(**input)来代替即可。至于为什么咱也不清楚😭。
同样还有其他方法——使用no_grad()方法：
# validate the model
if gpu==0 :
  with torch.no_grad():
    model.eval()
    for data, target in valid_loader:
      if torch.cuda.is_available:
        data, target = data.cuda(), target.cuda()
        output = model(data)
        loss = criterion(output, target)
        valid_loss += loss.item()*data.size(0)

模型保存

模型的保存与加载，与单GPU的方式有所不同。这里通通将参数以cpu的方式save进存储, 因为如果是保存的GPU上参数，pth文件中会记录参数属于的GPU号，则加载时会加载到相应的GPU上，这样就会导致如果你GPU数目不够时会在加载模型时报错，像下面这样：

RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.

模型保存都是一致的，不过时刻记住方案二中你有多个进程在同时跑，所以会保存多个模型到存储上，如果使用共享存储就要注意文件名的问题，当然一般只在rank0进程上保存参数即可，因为所有进程的模型参数是同步的。

torch.save(model.module.cpu().state_dict(), "model.pth")

模型的加载：

param=torch.load("model.pth")

以下是 huggingface/transformers代码中用到的模型保存代码

if torch.distributed.get_rank() == 0:
    model_to_save = model.module if hasattr(model, "module") else model  # Take care of distributed/parallel training（因为并行训练的方法，会在模型的key中多‘module’，也不知道为什么）
    model_to_save.save_pretrained(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

同一台机器上跑多个 ddp task

假设想在一台有4核GPU的电脑上跑两个ddp task，每个task使用两个核，很可能会需要如下错误：

RuntimeError: Address already in use
RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1544081127912/work/torch/lib/c10d/ProcessGroupNCCL.cpp:260, unhandled system error

原因是两个ddp task通讯地址冲突，这时候需要显示地设置每个task的地址

specifying a different master_addr and master_port in torch.distributed.launch

# 第一个task
export CUDA_VISIBLE_DEVICES="0,1" 
python -m torch.distributed.launch --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29501 train.py

# 第二个task
export CUDA_VISIBLE_DEVICES="2,3" 
python -m torch.distributed.launch --nproc_per_node=2 --master_addr=127.0.0.2 --master_port=29502 train.py

分布式模型部署，引入SyncBN，将普通BN替换成SyncBN(NLP中一般是没有BN的)。

为什么使用 SyncBN看这里：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)

参考

文章出处登录后可见！

已经登录？立即刷新

PyTorch 单机多GPU 训练方法与原理整理

相关推荐