PyTorch 分布式训练方法

Vachel    June 18, 2021

image.png

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较慢,且显卡资源利用度不高,不推荐使用多机多卡带来的加速效果非常显著。整体结果如下:

result.png

nn.DataParallel

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

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

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

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

image.png

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

image.png

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

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

image.png

最终,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之间通信使用的后端和端口:

image.png

然后,使用 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 进行训练: image.png

值得注意的是**,与****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)。

image.png

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

image.png

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

image.png

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

apex

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

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

image.png

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

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

image.png

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

image.png

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

image.png

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

多机多卡训练加速

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

分布式评价

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

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

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

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

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

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

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

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

image.png

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

总结

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


本文作者: Vachel


您可以转载、不违背作品原意地摘录及引用本技术博客的内容,但必须遵守以下条款: 署名 — 您应当署名原作者,但不得以任何方式暗示幻方为您背书,亦不会对幻方的权利造成任何负面影响。 非商业性使用 — 您不得将本技术博客内容用于商业目的。 禁止演绎 — 如果基于该内容改编、转换、或者再创作,您不得公开或分发被修改内容,该内容仅可供个人使用。