1. 为什么需要分布式训练
随着人工智能与深度学习的发展，大规模和超大规模的模型越来越受到业界的推崇。以NLP行业为例，从最开始的Bert-base只有1亿左右的参数量，到千亿级别的GPT-3，再到今年6月发布的目前全球最大预训练模型“悟道2.0”，参数规模达到惊人的1.75万亿，整个业界都由一种向更大模型发展的趋势。面对如此庞大的模型，必然也需要庞大的数据量才能进行训练，如果没有分布式训练的大算力加持，一个Epoch可能就要训练到天荒地老。抛开业界淬炼超大模型的场景，对于一个AI行业的普通算法工程师，面对日常的工作，分布式训练也可以大大加速模型的训练、调参的节奏、以及版本的迭代更新，在时间如此珍贵的当下，相信没有工程师会抗拒分布式训练带来的收益。因此，我们今天就聊聊深度学习中关于分布式训练的那些事儿。
2. 分布式训练策略
分布式训练策略按照并行方式不同，可以简单的分为数据并行和模型并行两种方式。
2.1 数据并行
数据并行是指在不同的GPU上都copy保存一份模型的副本，然后将不同的数据分配到不同的GPU上进行计算，最后将所有GPU计算的结果进行合并，从而达到加速模型训练的目的。由于数据并行会涉及到把不同GPU的计算结果进行合并然后再更新模型，根据跟新方式不同，又可以分为同步更新和异步更新。在数据并行中，每个GPU只计算一个batch中的一部分数据，同步更新指的就是在等待所有的GPU都计算完成之后，然后再统一合并和更新网络的权重，并广播到所有的GPU中，随后进行下一轮的计算。而异步跟新不同，异步更新中每个GPU在独立计算完成之后，都无需等待其他GPU，可以立即更新整体权重，然后广播到其他GPU中，随后马上进入下一轮的计算。由此可见，同步更新需要等待所有的GPU都计算完成才能更新，如果集群中某一个GPU训练慢了，或者集群中的通信出现抖动，都会影响到整个网络的训练速度，类似木桶效应，最短板决定了最大的容量。而异步更新由于不用等待其他GPU节点，因此总体训练速度会快一些，但是会有一个严重的梯度失效的问题。即在异步的情况下，每一个节点完成训练之后，都会马上去更新，这会造成其他节点现在的模型参数和这一轮训练前采用的模型参数可能不一致，从而导致此时的梯度过期。因此，异步更新虽然快，但是由于梯度失效问题，模型往往会陷入到次优解中。
2.2 模型并行
与数据并行不同，分布式训练中的模型并行是指将整个神经网络模型拆解分布到不同的GPU中，不同的GPU负责计算网络模型中的不同部分。这通常是在网络模型很大很大、单个GPU的显存已经完全装不下整体网络的情况下才会采用。由于深度学习的模型通常包含很多层，层与层之间的运行有先后训练，前向传播和反向梯度计算的时候，前面的层和后面的层都会彼此依赖作为输入输出，因此这种串行的逻辑对加速造成了一定的限制。但是相比起来，我们也算可以通过模型并行的方式把一个超大模型训练起来，不然对于一个单GPU的话，超大模型是完全没办法work的。 因此，对比起来，模型并行由于各个GPU只加载了模型的一部分网络结构，存在一定的依赖关系，造成了规模的伸缩性比较差，不能随意的增减GPU的数量，因此在实际中运用的并不多。而数据并行的方式，由于各个GPU相互独立，方便GPU的扩缩容，同时加速效果好，因此在实际中运用较多，但是在某些时候，我们也可以同时结合数据并行和模型并行两种方式。
3. 基于Pytorch的分布式训练方法
在Pytorch中为我们提供了两种多GPU的分布式训练方案：torch.nn.DataParallel（DP）和 torch.nn.parallel.Distributed Data Parallel（DDP）。
3.1 Data Parallel
DP模式使用起来非常容易，只需要对单GPU的代码修改其中一行就可以运行了，由于DP模式采用的是PS架构，存在负载不均衡问题，主卡往往会成为训练的瓶颈，因此训练速度会比DDP模式慢一些。而且DP只支持单机多卡的方式，一般一台机器只能安装最多8张卡，当我们要训练特别大型的任务时，8卡就会显得特别吃紧，因此会有一定的限制。
# use DataParallel
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = torch.nn.DataParallel(model)
model.to(device)
3.2 DistributedDataParallel
与DP模式不同，DDP模式本身是为多机多卡设计的，当然在单机多卡的情况下也可以使用。DDP采用的是all-reduce架构，基本解决了PS架构中通信成本与GPU的数量线性相关的问题。虽然在单机多卡情况下，可以使用DP模式，但是使用DDP通常会比DP模式快一些，因此DDP模式也是官方推荐大家使用的方式。改造现有的代码使用DDP也非常方便，通过下面几个步骤就可以轻松搞定。
# 1. init backend nccl
torch.distributed.init_process_group(backend='nccl')
# 2. config gpu
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
# 3. use DistributedSampler
training_loader = DataLoader(training_set, batch_size=TRAIN_BATCH_SIZE, sampler=DistributedSampler(training_set))
# 4. move model to gpu
model.to(device)
# 5. use DistributedDataParallel
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
3.3 Horovod
除了Pytorch原生提供的DP和DDP方式以外，也有很多优秀的由第三方提供的分布式训练工具，其中Horovod就是比较常用的一款。Horovod是Uber开源的跨平台分布式训练框架（horovod名字来源于俄罗斯一种民间舞蹈，舞者手拉手站成一个圆圈跳舞，类比了GPU设备之间的通信模式，如果该框架是中国人或者华人开发的话，我估计可能就叫“锅庄”了吧^-^），从名字可以看出来，Horovod采用all-reduce架构来提高分布式设备的通信效率。同时，Horovod不仅支持Pytorch，也支持TensorFlow等其他深度学习框架。训练中如果想使用Horovod的话，其实对代码的改动也比较少，如下所示。
import horovod.torch as hvd
# 1. init horovod
hvd.init()
# 2. Pin GPU to be used to process local rank (one GPU per process)
torch.cuda.set_device(hvd.local_rank())
# 3. Partition dataset among workers using DistributedSampler
train_sampler = DistributedSampler(training_set, num_replicas=hvd.size(), rank=hvd.rank())
training_loader = DataLoader(training_set, batch_size=TRAIN_BATCH_SIZE, sampler=train_sampler)
# 4. Add Horovod Distributed Optimizer
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
# 5. Horovod: broadcast parameters from rank 0 to all other processes.
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
另外，字节跳动也开源了一款高性能的分布式深度学习训练框架BytePS（项目github地址：https://github.com/bytedance/byteps），该框架没有采用热门的all-reduce，反而采用了PS架构，通过利用额外的CPU资源作为Parameter Server等措施，提升了通信性能，据说效果可以优于Horovod。而在几天之前，快手联合苏黎世理工也宣布开源了一款分布式训练框架Bagua（八卦），Bagua专门针对分布式场景设计了特定的优化算法，实现了算法和系统层面的联合优化，性能较同类提升60%。感兴趣的同学也可以关注一下。项目github地址：https://github.com/BaguaSys/bagua
4. 实验对比
这里我们对比了Pytorch原生的DP和DDP模式，同时也选择了第三方插件Horovod进行对比。实验选择了基于bert-base的预训练语言模型进行文本分类的任务。具体实验参数如下：GPU型号: V100, learning_rate: 2e-5, batch_size: 128, max_len: 128, epochs: 1, train_set_size: 48w 由于DDP和Horovod都是采用all-reduce架构，因此性能相当，可见Pytorch原生的DDP模式也已经做得非常不错了。而DP相比其他模式性能就会差一些。因此在实际工作中，还是比较推荐使用DDP或者Horovod进行分布式训练。
总结
本文探讨了深度学习中模型并行和数据并行的分布式策略，并基于Pytorch框架介绍了原生的DP和DDP模式，以及第三方Horovod分布式训练框架。从后面的实验对比可以看出，平时工作中比较推荐使用DDP或者Horovod的方式。分布式训练是深度学习中非常重要的一环，除了Horovod，其他各大厂商也相继开源了自己的分布式训练框架，比如BytePS、DeepSpeed、Bagua等等，这些框架的开源也将进一步推动这个领域的发展，为深度学习提供更优秀的工具。
作者简介
Hongyu OPPO高级NLP算法工程师
主要从事NLP、知识图谱及相关领域的工作