PyTorch 分布式训练方法

2018年，将近3亿参数的Bert模型横空出世，将NLP领域推向了新的高度。近年来人工智能领域的发展愈来愈趋向于对大模型的研究，各大AI巨头都纷纷发布了其拥有数千亿参数的大模型，诞生出了很多新的AI应用场景。 ​ 另一方面，多种因素继续推动大模型的长足发展：1) 社会正经历着深度的数字化转型，大量的数据逐渐融合，催生出许多人工智能的应用场景和需求；2) 硬件技术不断进步：英伟达 A100 GPU，Google的 TPU，阿里的含光800等，推动着AI算力的不断提升。可以这么说，“大数据+大模型+大算力”的组合，是迈入 AI 2.0 的基石。 ​ 然而，对于大多数AI从业者们来说，受限于一些因素，平常很难接触到大数据或者大算力，很多高校研究生们往往都习惯于用一张显卡（个人PC）来加速训练，这不仅训练缓慢，甚至会因此限制住了想象，禁锢了创造力。 ​ 幻方AI根据自身业务需求，构建了“萤火二号”，其旨在提供澎湃如海的深度学习算力服务。团队自研了智能分时调度系统、高效存储系统与网络通信系统，可以将集群当成一台普通计算机来使用，根据任务需求弹性扩展GPU算力；自研了hfai数据仓库、模型仓库，优化AI框架与算子，集成落地了许多前沿的应用场景；通过Client接口或者Jupyter就可以轻松接入，加速训练起AI大模型。

本期文章分享的，是如何使用起多张显卡，来加速你的AI模型。分布式训练技术逐渐成为AI从业者必备技能之一，这是从“小模型”走向“大模型”的必由之路。我们以 PyTorch 编写的ResNet训练为例，为大家展示不同的分布式训练方法及其效果。

训练任务：ResNet + ImageNet

训练框架：PyTorch 1.8.0

训练平台：幻方萤火二号

训练代码：https://github.com/HFAiLab/pytorch_distributed

训练准备

萤火二号当前完整支持了PyTorch的并行训练环境。笔者先使用1个节点8张卡，测试不同的分布式训练方法所需要的时间、显卡利用效率；之后采用多个节点，来验证并行加速的效果。测试的方法如下：

1. nn.DataParallel
2. torch.distributed + torch.multiprocessing
3. apex 半精度

为充分利用显存，这里batch_size设为400，记录每Epoch的耗时来进行对比。我们先采用单机单卡测试作为baseline，每Epoch耗时1786秒上下。 ​

测试结果发现，有半精度加持的并行加速效果最好；DataParallel较慢，且显卡资源利用度不高，不推荐使用；多机多卡带来的加速效果非常显著。整体结果如下：

nn.DataParallel

DataParallel 是PyTorch早期提出的数据并行训练方法，其使用单进程控制，将模型和数据加载到多个GPU中，控制数据在 GPU 之间的流动，协同不同 GPU 上的模型进行并行训练。 ​

使用起来非常方便，我们只需要用 nn.DataParallel 包装模型，再设置一些参数即可。需要定义的参数包括：

• 参与训练的 GPU 有哪些，device_ids=gpus；
• 用于汇总梯度的 GPU 是哪个，output_device=gpus[0] 。

DataParallel 会自动帮我们将数据切分 load 到相应 GPU，将模型复制到相应 GPU，进行正向传播计算梯度并汇总：

值得注意的是，这里因为有8张显卡，**batch_size**要设为3200，由DataParallel均分到不同的显卡里。模型和数据都要先 load 进 GPU中，DataParallel 的 module 才能对其进行处理，否则会报错。

发起训练。从下图我们可以看到，在显存几乎占满的情况下，8张显卡平均GPU利用率不高（36.5%）。

​ 单独查看每一张显卡的利用情况，可以发现资源使用不平均。其中0号显卡作为梯度汇总，资源利用度相比其他显卡要高一点，整体利用率不到50%。

最终，nn.DataParrallel方法下，ResNet平均每Epoch耗时984s左右，相比单机单卡加速大约两倍效果。 ​

torch.distributed

在 pytorch 1.0 之后，PyTorch官方对分布式的常用方法进行了封装，支持 all-reduce，broadcast，send 和 receive 等等。通过 MPI 实现 CPU 通信，通过 NCCL 实现 GPU 通信，解决了 DataParallel 速度慢，GPU 负载不均衡的问题。 ​

与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，我们只需要编写一份代码，torch 就会自动将其分配给 n 个进程，分别在 n 个 GPU 上运行。 ​

具体的，首先使用 init_process_group 设置GPU之间通信使用的后端和端口：

​

然后，使用 DistributedDataParallel 包装模型，它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值：

model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[local_rank])

​

最后，因为是采用了多进程管理，需要使用 DistributedSampler 对数据集进行划分，它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练：

值得注意的是**，与****nn.DataParrallel**不同，这里**batch_size**应该为400，因为其只负责当前rank下对应的partition，8张卡一共组成3200个样本的batch_size。 ​

至于 API 层面，PyTorch 为我们提供了 torch.distributed.launch 启动器，用于在命令行分布式地执行 python 文件。在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 通过参数传递给 python，我们可以这样获得当前进程的 index：

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
args = parser.parse_args()
print(args.local_rank)

启动8个命令行，执行如下命令：

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8 train.py

torch.multiprocessing ​

手动启动命令行比较麻烦，这里有一种更简单的方法：torch.multiprocessing 会帮助自动创建进程，绕开 torch.distributed.launch 自动控制开启和退出进程的一些小毛病。 ​

如下面的代码所示，将原本需要 torch.distributed.launch 管理的执行内容封装进 main 函数中，通过 torch.multiprocessing.spawn 开启了 nprocs=8 个进程，每个进程执行 main 并向其中传入 local_rank（当前进程 index）。

发起训练。从下图我们可以看到，显存几乎占满，8张显卡平均利用率高（95.8%）。

​

同时，每张显卡利用都比较充分，利用率都在80%以上：

​

最终，torch.distributed方法下，ResNet平均每Epoch耗时239s左右，相比单机单卡加速大约八倍效果。 ​

apex

Apex 是 NVIDIA 开源的用于混合精度训练和分布式训练库。Apex 对混合精度训练的过程进行了封装，改两三行配置就可以进行混合精度的训练，从而大幅度降低显存占用，节约运算时间。此外，Apex 也提供了对分布式训练的封装，针对 NVIDIA 的 NCCL 通信库进行了优化。PyTorch 1.6版本之后，开始原生支持，即torch.cuda.amp。 ​

直接使用 amp.initialize 包装模型和优化器，apex 就会自动帮助我们管理模型参数和优化器的精度了，根据精度需求不同可以传入其他配置参数。

在分布式训练的封装上，Apex 的改动并不大，主要是优化了 NCCL 的通信。因此，大部分代码仍与 torch.distributed 保持一致。使用的时候只需要将 torch.nn.parallel.DistributedDataParallel 替换为 apex.parallel.DistributedDataParallel 用于包装模型。 ​

在正向传播计算 loss 时，Apex 需要使用 amp.scale_loss 包装，用于根据 loss 值自动对精度进行缩放：

发起训练。从下图我们可以看到，在同等batch_szie条件下，显存只占用了60%，而8张显卡平均利用率高（95.8%）。

​

同时，每张显卡利用都比较充分，利用率都在80%以上：

最终，apex方法下，ResNet平均每Epoch耗时230s左右，相比torch.distributed加速了一点。同时，apex也节约了GPU的算力资源，对应的可以设置更大的batch_size，获取更快的速度。我们极限测试下，能在torch.distributed基础上再提升50%的性能左右。 ​

多机多卡训练加速

从上面的试验可以看到，8张卡带来的加速效果近似8倍。如果我们使用更多的机器，是否会有更快的提速效果呢？这里，我们对 apex的代码，配置4个节点，一共32张A100 GPU来进行试验，获得了惊人的效果。 ​ ​ 可以看到，每Epoch总耗时52秒左右，相比单机单卡速度提升30多倍，几乎和显卡数量保持一致。 ​

分布式评价

以上展示了分布式训练的过程。然而，如何对训练的结果进行快速的推理与测试，也是非常重要的问题，例如：

1. 训练样本被切分成了若干个部分，被若干个进程分别控制运行在若干个 GPU 上，如何在进程间进行通信汇总这些（GPU 上的）信息？

2. 使用一张卡进行推理、测试太慢了，如何使用 Distributed 进行分布式地推理和测试，并将结果汇总在一起？

3. …

要解决这些问题，我们需要一个更为基础的 API，汇总记录不同 GPU 上生成的准确率、损失函数等指标信息。这个 API 就是 torch.distributed.reduce。

reduce等 API 是 torch.distributed 中更为基础和底层的 API。这些 API 可以帮助我们控制进程之间的交互，控制 GPU 数据的传输。在自定义 GPU 协作逻辑，汇总 GPU 间少量的统计信息时，大有用处。熟练掌握这些 API 也可以帮助我们自己设计、优化分布式训练、测试流程。 ​

如上图所示，它的工作过程包含以下两步：

1. 在调用 reduce(tensor, op=…)后，当前进程接受其他进程发来的 tensor。
2. 在全部接收完成后，当前进程（例如rank 0）会对当前进程的和接收到的 tensor 进行 op 操作。

通过上述步骤，我们就能够对不同 GPU 上的训练数据的损失函数进行求和了：

除了reduce，PyTorch官网提供了诸如 scatter, gather, all-reduce等6种集合通信方案，具体的参考官方文档： https://ptorch.com/docs/1/distributed ​

总结

本文介绍了不同的PyTorch并行训练方法，并在幻方萤火二号上进行了测试。从测试结果中我们能够得出，多机多卡并行训练能有效提升我们效率，apex 方法结合torch.distributed 能最大限度的发挥显卡的算力，让加速效果与显卡数量一致。同时，我们也发现，随着显卡数量的增多，受限于梯度汇合中的reduce计算，单个显卡的利用率会逐渐降低。如何高效优化reduce算法，尽可能提高gpu算力，从而加速整体训练效率，是接下来值得大家继续研究的课题。