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

模型的多GPU并行训练，DDP

article 2024/12/25 16:22:23

DDP全称是DistributedDataParallel, 在torch.nn.parallel里面。
今天总结一下用DDP进行多GPU并行训练的方法，
内容来自build gpt2加上自己的补充。

如果你有多块GPU，就可以充分利用它们。
DDP会创建多个process（进程，不是线程哦）, 每个process分配一个GPU，这些process同时运行，于是就达到了多个GPU同时训练的目的。
之前写的训练代码你就可以想象它们同时在被多个process同时运行，
但每个process处理的是不同部分的数据（数据当然要不同了，不然不成了重复训练了嘛）。
这时候你还需要注意一个问题就是梯度的计算，比如现在有8个GPU，它们同时计算了不同部分数据的梯度，
而你的训练目标是不是所有数据汇总的梯度，你就需要把8个GPU计算的梯度加起来求平均。

DDP变量

那么你该如何区分哪个机器上的哪个进程呢，有下面几个变量帮你区分：
WORLD_SIZE: 总进程数
LOCAL_RANK: 当前进程在本机器上的局部排名,从0开始计数
RANK: 当前进程在所有进程中的全局排名,从0到 WORLD_SIZE-1
举个例子：
在2台机器上运行4个进程时:
机器1:

进程0: RANK=0, LOCAL_RANK=0
进程1: RANK=1, LOCAL_RANK=1

机器2:

进程2: RANK=2, LOCAL_RANK=0
进程3: RANK=3, LOCAL_RANK=1

所有进程的 WORLD_SIZE=4
如果只有一个机器，那local_rank和rank的值是相同的。

启动train.py

可以手动设置分布式环境

import torch.distributed as dist

# 手动设置环境变量
os.environ['RANK'] = '0'
os.environ['WORLD_SIZE'] = '4'
os.environ['LOCAL_RANK'] = '0'
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'

# 初始化进程组
dist.init_process_group(backend='nccl')

更方便的是用torchrun命令，它可以自动做如下事情：

自动设置环境变量(RANK, LOCAL_RANK, WORLD_SIZE等)
为每个GPU创建一个进程
设置进程组通信,初始化分布式环境
在每个进程中运行train.py

# 单机4卡
torchrun --nproc_per_node=4 train.py

# 多机训练
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=29500 train.py

torchrun会根据GPU序号设置RANK。例如在单机4卡时：

GPU 0: RANK=0
GPU 1: RANK=1
GPU 2: RANK=2
GPU 3: RANK=3

在多机时，RANK会跨机器累加。例如2机4卡：

机器1: RANK=0,1,2,3
机器2: RANK=4,5,6,7

那你说，我的train.py有时会用单个GPU，有时会用多个，我想让它具有通用性，不需要每次都修改代码。
可以下面这样做，后面都说的是单个机器的情况。

GPU之间的通信一般用’nccl’,
NCCL (NVIDIA Collective Communications Library) 是NVIDIA开发的GPU通信库，专门优化了GPU之间的通信效率。在PyTorch分布式训练中，它是GPU训练的默认后端选项。

from torch.distributed import init_process_group
    #如果用torchrun启动，RANK的值在0~WORLD_SIZE-1,WORLD_SIZE为进程数，即GPU总数
    ddp = int(os.environ.get('RANK', -1)) != -1
    if ddp:
        assert torch.cuda.is_available() #DDP中CUDA是必须的
        init_process_group(backend='nccl')
        ddp_rank = int(os.environ['RANK'])
        ddp_local_rank = int(os.environ['LOCAL_RANK'])
        ddp_world_size = int(os.environ['WORLD_SIZE'])
        device = f'cuda:{ddp_local_rank}'
        torch.cuda.set_device(device)
        master_process = ddp_rank==0 #防止log重复输出,指定master process输出
    else:
        ddp_rank = 0
        ddp_local_rank = 0
        ddp_world_size = 1
        master_process = True
        device = 'cpu'
        if torch.cuda.is_available():
            device = 'cuda'
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            device = 'mps'  #较新的MAC用
        print(f'using device: {device}')