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

模型的多GPU并行训练，DDP

news 来源：原创 2025/5/20 20:01:14

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)) != -1if 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 = 0ddp_local_rank = 0ddp_world_size = 1master_process = Truedevice = '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}')

batch size超过GPU内存的情况

比如gpt2, paper中的batch size为0.5M个token, 一般取2的N次方会利于GPU加速计算，
所以令total_batch_size=2¹⁹=524288, 这么大的batch size一般GPU很难容纳，如何在有限的GPU内存上实现如此大的batch size?

可以每次处理小批量的batch, 分多次处理。
比如gpt的sequence length设为T=1024，batch size B=16, 这是一次处理的batch量，
总共要处理steps次，steps=total_batch_size // (B * T)

待更…

DDP变量

启动train.py

batch size超过GPU内存的情况

划分不同数据区域

汇总梯度计算

输出log

相关文章：