实验｜Vachel 算力支持｜幻方AIHPC
2018年，将近3亿参数的Bert模型横空出世，将NLP领域推向了新的高度。近年来人工智能领域的发展愈来愈趋向于对大模型的研究，各大AI巨头都纷纷发布了其拥有数千亿参数的大模型，诞生出了很多新的AI应用场景。
另一方面，多种因素继续推动大模型的长足发展：1)社会正经历着深度的数字化转型，大量的数据逐渐融合，催生出许多人工智能的应用场景和需求；2) 硬件技术不断进步：英伟达 A100GPU，Google的 TPU，阿里的含光800等，推动着AI算力的不断提升。可以这么说，“大数据+大模型+大算力”的组合，是迈入 AI 2.0的基石。
然而，对于大多数AI从业者们来说，受限于一些因素，平常很难接触到大数据或者大算力，很多高校研究生们往往都习惯于用一张显卡（个人PC）来加速训练，这不仅训练缓慢，甚至会因此限制住了想象，禁锢了创造力。
幻方AI根据自身业务需求，构建起了拥有 5000 张 A100 GPU的超级计算机“萤火二号”，其旨在提供澎湃如海的深度学习算力服务。团队自研了智能分时调度系统、高效存储系统与网络通信系统，可以将集群当成一台普通计算机来使用，根据任务需求弹性扩展GPU算力；自研了hfai数据仓库、模型仓库，优化AI框架与算子，集成落地了许多前沿的应用场景；通过Client接口或者Jupyter就可以轻松接入，加速训练起AI大模型。幻方AI旨在打造一个大规模AI计算服务，助力数据科学家与AI开发者们打破禁锢，推动AI新发展。
本期文章分享的，是如何使用起多张显卡，来加速你的AI模型。分布式训练技术逐渐成为AI从业者必备技能之一，这是从“小模型”走向“大模型”的必由之路。我们以 PyTorch 编写的ResNet训练为例，为大家展示不同的分布式训练方法及其效果。
训练任务：ResNet + ImageNet
训练框架：PyTorch 1.8.0
训练平台：幻方AI HPC
训练代码：https://github.com/HFAiLab/pytorch_distributed
训练准备
萤火二号当前一共有625个节点，每个节点8张A100，完整支持了PyTorch的并行训练环境。笔者先使用1个节点8张卡，测试不同的分布式训练方法所需要的时间、显卡利用效率；之后采用多个节点，来验证并行加速的效果。测试的采用三种方法：
1 nn.DataParallel
2 torch.distributed + torch.multiprocessing
3 apex
为充分利用显存，这里 batch_size设为400，记录每Epoch的耗时来进行对比。我们先采用单机单卡测试作为baseline，每Epoch耗时1786秒上下。
测试结果发现，有apex加持的并行加速效果最好；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+torch.multiprocessing
在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多倍，几乎和显卡数量保持一致。
分布式评价
以上展示了分布式训练的过程。然而，如何对训练的结果进行快速的推理与测试，也是非常重要的问题，例如：
训练样本被切分成了若干个部分，被若干个进程分别控制运行在若干个 GPU 上，如何在进程间进行通信汇总这些（GPU 上的）信息？
使用一张卡进行推理、测试太慢了，如何使用 Distributed 进行分布式地推理和测试，并将结果汇总在一起？
......
要解决这些问题，我们需要一个更为基础的 API，汇总记录不同 GPU 上生成的准确率、损失函数等指标信息。这个 API 就是torch.distributed.reduce。
reduce等 API 是 torch.distributed中更为基础和底层的 API。这些 API 可以帮助我们控制进程之间的交互，控制 GPU 数据的传输。在自定义 GPU 协作逻辑，汇总 GPU 间少量的统计信息时，大有用处。熟练掌握这些 API 也可以帮助我们自己设计、优化分布式训练、测试流程。
如上图所示，它的工作过程包含以下两步：
在调用reduce(tensor, op=...)后，当前进程接受其他进程发来的tensor。
在全部接收完成后，当前进程（例如rank 0）会对当前进程的和接收到的tensor进行 op 操作。
通过上述步骤，我们就能够对不同 GPU 上的训练数据的损失函数进行求和了：
除了reduce，PyTorch官网提供了诸如scatter,gather,all-reduce等6种集合通信方案，具体的参考官方文档：https://ptorch.com/docs/1/distributed
总结
本文介绍了不同的PyTorch并行训练方法，并在幻方AI
HPC上进行了测试。从测试结果中我们能够得出，多机多卡并行训练能有效提升我们效率，Apex 方法结合torch.distributed
能最大限度的发挥显卡的算力，让加速效果与显卡数量一致。同时，我们也发现，随着显卡数量的增多，受限于梯度汇合中的reduce计算，单个显卡的利用率会逐渐降低。如何高效优化reduce算法，尽可能提高gpu算力，从而加速整体训练效率，是接下来值得大家继续研究的课题。
参考文档：《当代研究生应当掌握的并行训练方法（单机多卡）》
https://zhuanlan.zhihu.com/p/98535650
幻方AI团队， 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI实验落地与价值创造，降低超大算力需求的使用门槛，欢迎各方数据科学家与开发者们一同共建。
微信公众号 | HFAILAB
算力资源随时待命