Pytorch 使用多块GPU训练模型

2022-10-22 12:58:33

1. 先确定几个概念:
        ①分布式、并行:分布式是指多台服务器的多块GPU(多机多卡),而并行一般指的是一台服务器的多个GPU(单机多卡)。
        ②模型并行、数据并行:当模型很大,单张卡放不下时,需要将模型分成多个部分分别放到不同的卡上,每张卡输入的数据相同,这种方式叫做模型并行;而将不同的数据分配到不同的卡上,运行相同的模型,最后收集所有卡的运算结果来提高训练速度的方式叫数据并行。相比于模型并行,数据并行更为常用,以下我们主要讲述关于数据并行的内容。
        ③同步更新、异步更新:同步更新指所有的GPU都计算完梯度后,累加到一起求均值进行参数更新,再进行下一轮的计算;而异步更新指单个GPU计算完梯度后,无需等待其他更新,立即更新参数并同步。同步更新速度取决于最慢的那个GPU,异步更新没有等待,但是会出现loss异常抖动等问题,一般常用的是同步更新。
        ④group、world size、node、rank、local_rank:group指的是进程组,默认情况下只有一个主进程就只有一个组,即一个 world,当使用多进程时,一个 group 就有了多个 world;world size表示全局进程个数;node表示物理机器数量;rank表示进程序号;local_rank指进程内 GPU 编号。举个例子,三台机器,每台机器四张卡全部用上,那么有group=1,world size=12
        机器一:node=0   rank=0,1,2,3   local_rank=0,1,2,3      这里的node=0,rank=0的就是master
        机器二:node=1   rank=4,5,6,7   local_rank=0,1,2,3
        机器三:node=2   rank=8,9,10,11   local_rank=0,1,2,3

2.DP和DDP(pytorch使用多卡多方式)
        DP(DataParallel)模式是很早就出现的、单机多卡的、参数服务器架构的多卡训练模式。其只有一个进程,多个线程(受到GIL限制)。master节点相当于参数服务器,其向其他卡广播其参数;在梯度反向传播后,各卡将梯度集中到master节点,master节点收集各个卡的参数进行平均后更新参数,再将参数统一发送到其他卡上,参与训练的 GPU 参数device_ids=gpus;用于汇总梯度的 GPU 参数output_device=gpus[0]。DP的使用比较简单,需要修改的代码量也很少,在Pytorch中的用法如下:

from torch.nn import DataParallel

device = torch.device("cuda")
gpus = [0,1,2]
model = MyModel()
model = model.to(device)
model = DataParallel(model, device_ids=gpus, output_device=gpus[0])

        DDP(DistributedDataParallel)支持单机多卡分布式训练,也支持多机多卡分布式训练。相对于DP模式的最大区别是启动了多个进程进行并行训练。目前DDP模式只能在Linux下应用。其在Pytorch的用法如下:

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

#1.初始化group,使用默认backend(nccl)就行。如果是CPU模型运行,需要选择其他后端。
dist.init_process_group(backend='nccl')

#2.要加一个local_rank的参数
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1)
args = parser.parse_args()

#3.从外面得到local_rank参数,在调用DDP的时候,其会根据调用gpu自动给出这个参数,后面还会介绍。
#或者local_rank=torch.distributed.local_rank()
local_rank = args.local_rank

#4.根据local_rank来设定当前使用哪块GPU
torch.cuda.set_device(local_rank)

# 5.定义并把模型放置到单独的GPU上,在调用`model=DDP(model)`前。如果加载模型,也必须在这里做。
device = torch.device("cuda", local_rank)

#6.封装之前要把模型移到对应的gpu
model = Mymodel()
model.to(device)

#7.之后才是初始化DDP模型
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

        在使用时,调用 torch.distributed.launch 启动器启动。

python -m torch.distributed.launch --nproc_per_node=GPU数量 main.py

3.DP和DDP的优缺点
        DP的优点:使用起来非常简单,修改的代码量最少,只要像这样model = nn.DataParallel(model)包裹一下模型就行了。
        DP的缺点:速度慢,会造成负载不均衡的情况(master卡的显存可能使用更多),成为限制模型训练速度的瓶颈。只适用单机多卡,不适用多机多卡,DP使用单进程,性能不如DDP。
        DDP优点:使用多进程,训练速度大大提升,没有GIL contention;性能更优;模型广播只在初始化的时候,不在每次前向传播时,故训练加速。
        DDP缺点:代码改动较多,坑较多,需要试错攒经验。
        主要差异可以总结为以下几点:
        DDP支持模型并行,而DP并不支持,这意味如果模型太大单卡显存不足时只能使用前者;
        DP是单进程多线程的,只用于单机情况,而DDP是多进程的,适用于单机和多机情况,真正实现分布式训练;
        DDP的训练更高效,因为每个进程都是独立的Python解释器,避免GIL问题,而且通信成本低其训练速度更快,基本上DP已经被弃用;
        必须要说明的是DDP中每个进程都有独立的优化器,执行自己的更新过程,但是梯度通过通信传递到每个进程,所有执行的内容是相同的;

4.注意事项
        ①在用 DDP包裹模型之前需要先用模型送到device上,也就是要先送到GPU上,否则会报错:AssertionError: DistributedDataParallel device_ids and output_device arguments only work with single-device GPU modules, but got device_ids [1], output_device 1, and module parameters {device(type='cpu')};
        ②在开始使用DDP之前,需要用dist.init_process_group(backend='nccl', init_method='env://')初始化进程组,一般建议用nccl的backend,init_method不写默认就是’env://’。这个要写在最前面,不能运行两次,否则会报错:RuntimeError: trying to initialize the default process group twice!,在初始化进程组之后,经常会跟这样两句话:
        torch.cuda.set_device(args.local_rank)
        device = torch.device('cuda', args.local_rank)
        ③与单GPU模式下一个区别就是你需要用DistributedSampler包裹你的dataset得到sampler输入到dataloader里面,这时候在dataloader就不能指定shuffle这个参数了;
        ④batch size的区别:对于DP而言,输入到dataloader里面的batch_size参数指的是总的batch_size,例如batch_size=30,你有两块GPU,则每块GPU会吃15个sample;对于DDP而言,里面的batch_size参数指的却是每个GPU的batch_size,例如batch_size=30,你有两块GPU,则每块GPU会吃30个sample,一个batch总共就吃60个sample;
        ⑤load/save model的时候要注意一下,在multi-gpu下保存模型应该要用net.module.state_dict(),否则你在load的时候会有Missing key(s) in state_dict: "conv1.weight" ... Unexpected key(s) in state_dict: "module.conv1.weight",因为直接用net.state_dict()保存的模型会带有module的前缀,除非你自己又循环一遍参数,除去前缀,然后加载这个新的state_dict;
        ⑥打印信息,保存log,保存模型这些只要在用local_rank为0的进程进行就可以,因为模型会进行同步的,对于log信息,否则终端里就会显示很多快慢不一的信息,不美观,这也就是为什么很多代码里面都有args.local_rank==0的判断。

  • 作者:tiancanucas
  • 原文链接:https://blog.csdn.net/tiancanucas/article/details/124477063
    更新时间:2022-10-22 12:58:33