当前位置：首页 > news >正文

Pytorch分布式训练/多卡训练(二) —— Data Parallel并行(DDP)(2.1)(基本概念代码框架)

news 来源：原创 2024/4/27 22:00:43

Pytorch官网已经建议使用DistributedDataParallel来代替DataParallel, 因为DistributedDataParallel比DataParallel运行的更快, 然后显存分配的更加均衡. 而且DistributedDataParallel功能更加强悍

DDP通过Ring-Reduce的数据交换方法提高了通讯效率，并通过启动多个进程的方式减轻Python GIL的限制，从而提高训练速度。一般来说，DDP都是显著地比DP快，能达到略低于卡数的加速比(例如, 四卡下加速3倍)。现在几乎所有大牛的深度学习开源代码全都用DDP实现多卡训练。

Data Parallel的多卡训练的BN是只在单卡上算的，相当于减小了批量大小（batch-size）

DDP与DP的区别

①DDP在DataLoader部分需要使用Sampler，保证不同GPU卡处理独立的子集.

②DDP在模型部分使用DistributedDataParallel.

③DP不能做BN同步，DDP可以

④DP存在负载不均衡的问题(主卡的显存占用大于其他卡，甚至远大于。有时主卡爆了而其他卡的显存还没到一半)，而DDP没有

因为DP并不是完全的并行计算，只是数据在多张卡上并行计算，模型的保存和Loss的计算都是集中在几张卡中的一张上面，这也导致了多卡显存占用会不一致。

⑤DDP的速度要比DP更快

使用DistributedDataParallel，需要用torch.distributed.launch去launch程序

DDP多卡训练的原理

（1）将模型在各个GPU上复制一份；

（2）将总的batch数据等分到不同的GPU上进行计算(shuffle顺序打乱的)。每个进程都从磁盘加载其自己的数据

（3）在模型训练时，损失函数的前向传播和计算在每个 GPU 上独立执行。因此，不需要收集网络输出。在反向传播期间，各个进程通过一种叫Ring-Reduce的方法与其他进程通讯，交换各自的梯度，从而获得所有进程的平均梯度；然后用这个值在所有GPU上执行梯度下降

从而每个 GPU 在反向传播结束时最终得到平均梯度的相同副本。

（4）各个进程用平均后的梯度更新自己的参数，因为各个进程的初始参数、更新梯度是一致的，所以更新后的参数也是完全相同的。

这里的汇总还是由rank=0的卡汇总平均，然后再广播到其他卡。至于为什么不只用其他卡和主卡通信而是其他卡之间也要通信，是因为用的是ring-reduce, 组成一个环形来作信息传递，在多卡的情形下更高效

DDP利用All-Reduce，来进行不同进程上的梯度的平均操作(Ring-Reduce是All-Reduce的一个实现版本)

不同卡之间只有梯度在通信。

在Pytorch中使用DDP

DDP的官方最佳实践是，每个进程对应一张卡

举个例子：我有两台机子，每台8张显卡，那就是2x8=16个进程，并行数是16。

但是，我们也是可以给每个进程分配多张卡的。总的来说，分为以下三种情况：

每个进程一张卡。这是DDP的最佳使用方法。
每个进程多张卡，复制模式。一个模型复制在不同卡上面，每个进程都实质等同于DP模式。这样做是能跑得通的，但是，速度不如上一种方法，一般不采用。
每个进程多张卡，并行模式。一个模型的不同部分分布在不同的卡上面。例如，网络的前半部分在0号卡上，后半部分在1号卡上。这种场景，一般是因为我们的模型非常大，大到一张卡都塞不下batch size = 1的一个模型。

DDP的主要代码部分
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel

RANK = int(os.environ['SLURM_PROCID'])  # 进程序号，用于进程间通信
LOCAL_RANK = int(os.environ['SLURM_LOCALID']) # 本地设备序号，用于设备分配.
GPU_NUM = int(os.environ['SLURM_NTASKS'])     # 使用的 GPU 总数.
IP = os.environ['SLURM_STEP_NODELIST'] #进程节点 IP 信息.
BATCH_SIZE = 16  # 单张 GPU 的大小.

def dist_init(host_addr, rank, local_rank, world_size, port=23456):
    host_addr_full = 'tcp://' + host_addr + ':' + str(port)
    torch.distributed.init_process_group("nccl", init_method=host_addr_full,
                                         rank=rank, world_size=world_size)
    torch.cuda.set_device(local_rank)
    assert torch.distributed.is_initialized()

    
if __name__ == '__main__':
    dist_init(IP, RANK, LOCAL_RANK, GPU_NUM)
    
       # DataSet
    datasampler = DistributedSampler(dataset, num_replicas=GPU_NUM, rank=RANK)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=datasampler)

    # model 
    model = DistributedDataPrallel(model, 
                                   device_ids=[LOCAL_RANK], 
                                   output_device=LOCAL_RANK)
示例

未加入DDP的单GPU代码

运行: main.py
import torch
from torch import nn
from torch import optim

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = nn.Linear(10, 10).to(device)
# 前向传播
input = torch.randn(20, 10).to(device)
outputs = model(input)
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()
加入DDP的代码

使用torch.distributed.launch启动DDP模式

运行:
python -m torch.distributed.launch --nproc_per_node 4 main.py
使用torch.distributed.launch，就会自动给程序传入local_rank参数，所以我们必须要在程序里写上argparse接收
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch import nn
from torch import optim
import argparse

# 新增：从外面得到local_rank参数
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1)
FLAGS = parser.parse_args()
local_rank = int(FLAGS.local_rank)
 
# 新增：DDP backend初始化
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl')  # nccl是GPU设备上最快、最推荐的后端
 
# 构造模型
device = torch.device("cuda", local_rank)
model = nn.Linear(10, 10).to(device)
# 新增：构造DDP model
model = DDP(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()

DDP的两种多卡启动方式

① torch.distributed.launch

代码量更少，启动速度更快

首选这个
python -m torch.distributed.launch --help
-m是 run library module as a script

② torch.multiprocessing

拥有更好的控制和灵活性