pytorch分布式训练方法总结

0 概述

在深度学习中， 出于训练效率的考虑， 甚至有时候模型太大导致单个GPU卡放不下的情况， 这时候都需要用到分布式训练。 从大的方面分类， 并行训练可以分为数据并行， 模型并行以及混合并行3种。其中数据并行应用最为广泛， 也比较成熟。而模型并行目前还不够成熟， 缺乏统一的方案。本文主要介绍数据并行的方式， 并且主要关注pytorch训练框架。

pytorch的并行训练主要有3种方式:

• DP (DataParallel)
• DDP (Distributed DataParallel)
• Horovod

其中DP和DDP是pytorch原生支持的方式， Horovod是第三方支持。

1 DP

1.1 原理

DP有以下特点：

• DP基于PS架构， 只能用于单机多卡的场景， 无法满足多机的场景
• DP基于单进程多线程的方式实现
• DP的PS架构下需要一个节点作为server节点， 通常就是0号卡（非物理意义上的0号卡）， 有负载不均衡的问题

DP的具体过程是： 在前向传播过程中，输入的 batch 是平均分配到每个 device 中去，而网络模型需要拷贝到每个 device 中。在反向传播过程中，每个device会单独计算梯度， 并送到0号卡上进行累计， 并进行参数的更新， 更新完的参数会发给其他卡进行同步。

1.2 使用

DP的最大优点是使用非常简单， 只需要在原本单机版的代码中加上一行：

model = torch.nn.DataParallel(model)

使用DP时无须进行数据的划分， 数据会自动切分到不同的卡上， 每个卡上会分到batch_size/N 的数据， 程序中设定的batch_size为N个卡上的总和。

2 DDP

2.1 原理

DDP有以下特点：

• DDP基于 ring-all-reduce 进行通信， 能用于多机多卡的场景
• DDP基于多进程的方式实现
• DDP不需要中心节点， 负载均衡
• DDP的速度比DP快

DDP的具体过程是： 在每次迭代中，操作系统为每个GPU创建一个进程，每个进程具有自己的 optimizer ，并独立完成所有的优化步骤，进程内与一般的训练无异。在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。各进程用该梯度来更新参数。由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致。而在 DP中，全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。相较于DP ，DDP 传输的数据量更少，因此速度更快，效率更高。

2.2 使用

DDP的使用方式稍微复杂一些 ， 对照上图DDP的处理过程还是比较好理解的。

• 1 DDP 的初始化

# DDP：DDP backend初始化
torch.distributed.init_process_group(backend='nccl')  # nccl是GPU设备上最快、最推荐的后端
torch.cuda.set_device(local_rank)

注意： 一定要有torch.cuda.set_device(local_rank), 否则都会跑到一张卡上。
local_rank的获取有2种方式：

• 1. 通过 torch.distributed接口获取：

local_rank = torch.distributed.get_rank()
world_size = torch.distributed.get_world_size()

• 2. 通过命令行参数传递

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)

外部调用起进程的时候需要传入该参数， 如torch.distributed.launch方式启动时会自动传入该参数。

• 2 数据集的读取需要修改

dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True,transform=transform)
# 使用DistributedSampler  
sampler = torch.utils.data.distributed.DistributedSampler(dataset=dataset,
                                                          num_replicas=world_size,
                                                          rank=local_rank)
# 需要注意的是，这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
trainloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, sampler=sampler)

数据集的修改主要就是添加了分布式的sampler, 其中有2个很重要的参数：world_size和 rank, world_size就是总的卡数， 比如2机16卡， world_size就是16，rank_id就是0-15， 表示每个卡的编号。

• 3 模型构造

from torch.nn.parallel import DistributedDataParallel as DDP
model= model.cuda()
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

注意：model= model.cuda()要在调用DDP之前。

• 4 模型权重保存

if torch.distributed.get_rank() == 0:
   torch.save(model.module.state_dict(), "%d.ckpt" % epoch)

模型的保存有2点需要注意： 1）只用在rank=0的卡上保存就可以了， 避免重复保存；2）保存是model.module而不是model,因为model是被DDP封装起来的。

• 5 启动

启动方式有很多种:
方式1： 通过torch.distributed.launch启动
1个节点的启动方式：

python -m torch.distributed.launch --nproc_per_node=8  main.py

多个节点的启动方式: 假设一共有两台机器（节点1和节点2），每个节点上有8张卡，节点1的IP地址为192.168.1.1 占用的端口12355, 节点2的IP地址为192.168.1.2， 占用大的端口号为13356，启动的方式如下：

#节点1
python -m torch.distributed.launch --nproc_per_node=8
           --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
           --master_port=12355 main.py
#节点2
python -m torch.distributed.launch --nproc_per_node=8
           --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
           --master_port=12356 main.py

3 Horovord

Horovod 是Uber公司开发的深度学习工具， 支持多种框架的分布式训练。
有关Horovod的介绍和使用可参考官方文档：Horovod文档参考
这里主要总结pytorch如何使用horovod进行分布式训练。

• 1 Horovod初始化

hvd.init()

• 2 为每个GPU 分配一个进程

if torch.cuda.is_available():
    torch.cuda.set_device(hvd.local_rank())

• 3 对学习率进行放大
  一般来说， 增大batch_size时学习率也要进行对应的增大。

• 4 把优化器用Horovod进行包装

optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())

• 5 把0卡的初始参数广播到其他卡的进程

hvd.broadcast_parameters(model.state_dict(), root_rank=0)

4 参考

1 https://zhuanlan.zhihu.com/p/358974461 PyTorch分布式训练基础–DDP使用