pytorch单机多卡的简单原理与使用

Pytorch单机多卡训练

参考

1. Pytorch 单机多卡从入门到入土（坑点记录）
2. DDP系列（详细说明请戳）

DataParallel：只支持单机多卡，代码很方便，只需要添加一行，但是效率比较低，不推荐使用
DistributedDataParallel：支持多机多卡，效率高，但是要折腾一下代码，DDP是一个支持多机多卡、分布式训练的深度学习工程方法。

原理

通过提高batch size来增加并行度，Ring-Reduce的数据交换方法提高了通讯效率，并启动多个进程的方式减轻Python GIL的限制，从而提高训练速度。DDP都是显著地比DP快，能达到略低于卡数的加速比（例如，四卡下加速3倍）。
具体：

1. （缓解GIL限制）在DDP模式下，会有N个进程被启动，每个进程在一张卡上加载一个模型，这些模型的参数在数值上是相同的。
2. （Ring-Reduce加速）在模型训练时，各个进程通过一种叫Ring-Reduce的方法与其他进程通讯，交换各自的梯度，从而获得所有进程的梯度；
3. （实际上就是Data Parallelism）各个进程用平均后的梯度更新自己的参数，因为各个进程的初始参数、更新梯度是一致的，所以更新后的参数也是完全相同的。

对比

在DP模式中，总共只有一个进程（受到GIL很强限制）。master节点相当于参数服务器，其会向其他卡广播其参数；在梯度反向传播后，各卡将梯度集中到master节点，master节点对搜集来的参数进行平均后更新参数，再将参数统一发送到其他卡上。这种参数更新方式，会导致master节点的计算任务、通讯量很重，从而导致网络阻塞，降低训练速度。但是DP也有优点，优点就是代码实现简单。

Quick Demo

import os
# new
local_rank = int(os.environ["LOCAL_RANK"])
import torch
import torch.nn as nn
from torch import optim
import torch.distributed as dist
##############################################
# import argparse
# parser = argparse.ArgumentParser()
# parser.add_argument("--device", default=-1)
# FLAGS = parser.parse_args()
# local_rank = FLAGS.local_rank
##############################################

# 将模型加载到对应的gpu上
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl')
device = torch.device("cuda", local_rank)

# 构造模型
model = nn.Linear(10, 10).to(device)
model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)

# 前向传播
outputs = model(torch.randn(20, 10).to(device))
labels = torch.randn(20, 10).to(device)
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()
# 后向传播
optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.step()

如果使用注释中的代码并配合python -m torch.distributed.launch --nproc_per_node 4 main.py使用，也就是通过cmd读取local_rank的话将会报warning（如下），是因为这个方法将会被弃用。

FutureWarning: The module torch.distributed.launch is deprecated
and will be removed in future. Use torchrun.
Note that --use_env is set by default in torchrun.
If your script expects `--local_rank` argument to be set, please
change it to read from `os.environ['LOCAL_RANK']` instead. See 
https://pytorch.org/docs/stable/distributed.html#launch-utility for 
further instructions

取而代之，是代码前几行中注释了为new的代码，并且通过 torchrun main.py 直接运行即可。

使用

1. 每个进程一张卡，这是DDP的最佳使用方法。
2. 每个进程多张卡，并行模式。一个模型的不同部分分布在不同的卡上面。例如，网络的前半部分在0号卡上，后半部分在1号卡上。这种场景，一般是因为我们的模型非常大，大到一张卡都塞不下batch size = 1的一个模型。

基本概念

1. group
   即进程组。默认情况下，只有一个组。这个可以先不管，一直用默认的就行。
2. world size
   表示全局的并行数，如果只开了一个进程那就是1。

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

3. rank
   表现当前进程的序号，用于进程间通讯。对于16的world size来说，就是0,1,2,…,15。
   注意：rank=0的进程就是master进程。

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

3. 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()

增加Dataset Demo

## main.py文件
import os
import torch
import torch.nn as nn
from torch import optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
from torch.utils.data import DataLoader, data
import warnings 
warnings.filterwarnings('ignore')
local_rank = int(os.environ["LOCAL_RANK"])

torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl')
device = torch.device("cuda", local_rank)

# print(f"world size: {torch.distributed.get_world_size()}")
print(f"rank: {torch.distributed.get_rank()}")

class CPPDataset(data.Dataset):
    def __init__(self, seqs, labels):
        self.seqs = seqs
        self.labels = labels
    def __len__(self):
        return len(self.seqs)
    def __getitem__(self, idx):
        return self.seqs[idx], self.labels[idx]


# 构造模型
model = nn.Linear(10, 1).to(device)
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

X_train = torch.Tensor(64,10).to(device)
y_train = torch.zeros(64).to(device)
train_dataset = CPPDataset(X_train, y_train)

# 新增1：使用DistributedSampler，DDP帮我们把细节都封装起来了。用，就完事儿！
#       sampler的原理，后面也会介绍。
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
# 需要注意的是，这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=4, sampler=train_sampler)

loss_fn = nn.MSELoss()

for epoch in range(4):
    # 新增2：设置sampler的epoch，DistributedSampler需要这个来维持各个进程之间的相同随机数种子
    trainloader.sampler.set_epoch(epoch)
    # 后面这部分，则与原来完全一致了。
    for x, yy in trainloader:
        prediction = model(x)
        
        loss = loss_fn(prediction.squeeze(), yy)
        loss.backward()
        optimizer = optim.SGD(model.parameters(), lr=0.001)
        optimizer.step()
        
    # 保存模型参数只需要在rank=0上保存
    if dist.get_rank() == 0:
            checkpoint = {
                "net": model.module.state_dict()
            }
            torch.save(checkpoint, 'model.pt'))

最终通过运行torchrun --nproc_per_node 4 main.py 即可实现单机多卡训练。
当你希望在一个机器上同时跑两个torchrun程序，请使用torchrun --nproc_per_node 4 -master_port=22224 main.py，master_port的数字不能与正在运行的程序一致，会产生冲突。

注意点：

1. 保存模型 checkpoint、记录 tensorboard、输出准确率、甚至是 tqdm 都只让主进程（rank0）执行，其他进程不执行。
2. 在用 DDP 封装模型后，模型的本体（我们定义的那个类）是 model.module；所以保存模型时，最好保存 model.module.state_dict()，否则存下来的参数的 key 前面会多一个 module.，不便再次 load 模型。
3. 我们经常会在训练的每个 epoch 后进行 evaluate / inference，为避免有些进程测试完之后开始了下一轮训练，但其他进程还在测试，最好在每个 epoch 开始训练前（或者每个 epoch 完成训练后）用 dist.barrier() 同步一下。对于数据的读取是采用主进程预读取并缓存，然后其它进程从缓存中读取，不同进程之间的数据同步具体通过torch.distributed.barrier()实现。
4. 单卡 evaluate / inference 用的模型最好是本地模型 model.module 而非 DDP 包装的模型，否则非主进程会在第 3 条的 dist.barrier() 处卡死，推测原因是 DDP 包装的模型会在一些地方同步 bucket。
5. 如果模型在forward的时候出现没有使用的output，比如(logits, _)= model(x_pad, mask)，则有可能会报错，报错内容如下。需要在封装DDP的地方增加一个参数 find_unused_parameters=True)

RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss.