深度学习中的分布式训练

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 DataParallelif 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 nccltorch.distributed.init_process_group(backend='nccl')# 2. config gpulocal_rank = torch.distributed.get_rank()torch.cuda.set_device(local_rank)device = torch.device("cuda", local_rank)# 3. use DistributedSamplertraining_loader = DataLoader(training_set, batch_size=TRAIN_BATCH_SIZE, sampler=DistributedSampler(training_set))# 4. move model to gpumodel.to(device)# 5. use DistributedDataParallelif 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 horovodhvd.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 DistributedSamplertrain_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 Optimizeroptimizer = 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、知识图谱及相关领域的工作

获取更多精彩内容，搜索关注[OPPO 数智技术]公众号