pytorch中的分布式训练
pytorch分布式训练
·
1.简介
Pytorch的数据分为两种,torch.nn.parallel.DataParallel
(DP)和torch.nn.parallel.DistributedDataParallel
(DDP).
使用场景:
- DP模式用于单机多卡
- DDP模式可以用于单机多卡、多机多卡以及模型并行。
2. DP模式
DP模式用于单机多卡的情况。实现起来较为简单,只需要用torch.nn.parallel.DataParallel
包装模型即可,其余与单机单卡的情况相同。
1.1 流程
- master(一般是GPU0)从磁盘或者合页内存中取数据。
- master将数据分到其他GPU上
- master将模型复制到其他GPU上
- 每块GPU单独进行前向计算,得到输出
- master收集每块GPU上的输出,计算损失
- master将损失分到其他卡上,每块卡单独进行反向传播,计算梯度
- master收集每块GPU上的梯度,汇总以后,进行reduce操作,结果分发到每块卡上。
- 每块GPU根据梯度,单独更新模型
1.2 图示流程
可以看出DP模型比较多的操作是在0号卡上进行的。分数据、将模型的副本传到其他模型,收集每块卡的输出、计算loss,将loss传到每块卡上,在每块卡上进行反向传播得到梯度后,收集每块卡上的梯度,进行reduce上操作后,传到其他卡上。这样使得0号卡所占的内存比较多,使得内存使用不均衡,而且会经常出现其他卡等待0好卡计算的情形。
2. Distributed Data Parallel
2.1 处理流程
- 从master(一般是GPU0)从磁盘或者合页内存中取数据。
- 所有GPU同时去取数据,不需要GPU0去分
- 每块GPU单独进行前向计算
- 每块GPU单计算loss
- 每块GPU单独进行反向过程,计算出参数的梯度,并进行参数之间的传递(计算和参数传递间存在交叉过程)
- 在GPU0上进行梯度的allreduce操作,然后将梯度传递到其他GPU上。
- 每个GPU单独地进行参数更新
可以看出DDP只有在梯度收集的时候,存在GPU间的通信,其余的操作都是在各自的GPU上进行的,这样可以均衡负载,也可以节省时间。
1.3 实现
- 进程初始化
backend
可选nccl,gloo和mpi
模式,表示不同的collective communication方式。使用英伟达显卡,运行比较快的是nccl方式
local_rank= os.getenv('LOCAL_RANK', -1)
if local_rank != -1:
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
torch.distributed.init_process_group(backend="nccl", init_method='env://')
- 数据集创建
def create_dataloader(...):
train_datasets= torchvision.datasets.MNIST(
root='./mnist',
train=True,
transform=torchvision.transforms.ToTensor(), # (0,255) -->(0,1)
download=False
)
train_sampler = DistributedSampler(train_datasets) ## 防止不同进程交叉取数据
train_dataloader = DataLoader(train_datasets, sampler=train_sampler, batch_size=int(args.test_batch_size/args.world_size),
num_workers=args.num_workers, pin_memory=True)
...
....
train_sampler,train_dataloader,test_sampler,test_dataloader = create_dataloader(args,main_process)
- 使用DDP模型对模型进行包装
model = ...
model.to(device)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
print('use {} gpus!'.format(num_gpus))
model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
output_device=args.local_rank)
- 模型训练
for epoch in range(EPOCH):
# 步骤五:打乱数据顺序
train_sampler.set_epoch(epoch)
model.train()
start= time()
### train
for step, (b_x, b_y) in enumerate(train_dataloader):
predit_y = model(b_x)
loss = loss_func(predit_y, b_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
- 启动方式
- 单机多卡
python -m torch.distributed.run --nproc_per_node 2 --master_port 2393 ddp_train.py
- 多机多卡
在机器0上运行python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.0.0.1" --master_port=1234 ddp.py
在机器1上运行python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="192.0.0.1" --master_port=1234 ddp.py
不同机器上的node_rank
不同,其余相同
TODO
参考
[1] pytorch(分布式)数据并行个人实践总结——DataParallel/DistributedDataParallel
[2] PyTorch Distributed: Experiences on Accelerating Data Parallel Training
[3] 多机多卡
更多推荐
已为社区贡献1条内容
所有评论(0)