大规模并行分布式深度学习

作者：林伟

最新的十年是人工智能爆炸的十年。人工智能在图片、语音、语言理解、推荐算法等各个社会领域取得了重大突破，使其判断接近或超越了人类。这背后的主要原因是互联网的不断发展使我们能够快速积累大量的数据，加上硬件的快速发展和神经网络训练模式的创新，使我们越来越能够训练深而宽的神经网络，生产超越人类智能模型，并得到广泛的应用。因此，如何实现大规模并行分布式深度学习已成为深度学习的研究热点，成为促进算法创新的关键人工智能工程能力。

并行分布式深度学习的历史发展历程和范式

推广分布式训练的场景是大规模搜索广告推荐模式。随着互联网的发展，我们需要建立一个搜索广告推荐模型，可以满足互联网用户和数据规模。通过培训，我们需要有效地量化用户和网页，从而通过不同的人或网页的特征找到更多相关的信息，提高人们在海量数据中匹配数据的效率。这种模型通常需要学习一个非常大的特征矩阵，如现代系统的特征规模，从100亿[1]发展到1000亿甚至10000亿[2]，这需要一个分布式系统来支持这种机器学习模型训练。此外，由于人群和网页的特点，矩阵具有明显的热效应，非常稀疏，因此在访问时往往具有时间和空间的不均匀性。这也是 谷歌和其他团队在2014年 OSDI 发布分布式系统《Scaling Distributed Machine Learning with the Parameter Server》[3](其他相关研究也有[4][5])，该系统（Figure 1)建立分布式 Key-Value 存储系统保存了大量的特征向量，然后我们将大量的训练数据分发到多个无状态的训练节点（worker），每个工作节点从分布式开始 KV参数服务器获得最新参数，在当前数据培训计算估计和预期答案（通常是用户最终真实点击网页或商品作为目标答案）的差距，计算参数调整梯度，然后推动参数服务器相关加权平均，总结调整参数，使整个模型的预测更接近目标值。由于稀疏的特点，为了充分发挥系统的性能，这种模型往往追求动态目标，因为人们的兴趣和网页本身一直在变化。因此，我们可以利用这一特点进行异步训练，即在训练过程中适当使用过期参数和异步更新参数，从而更好地容忍参数服务器参数访问的不均匀性、分布式场景的容错性和网络通信的压力，从而提高分布式训练的效率。可以参考更具体 NIPS 2013 年的《More effective distributed ml via a stale synchronous parallel parameter server》等论文[5]。

同时，在图片、语音、语言等感知领域，随着大量数据的积累和计算能力的发展，神经网络算法终于可以取得突破性的结果 AlexNet[6], GoogleNet[7]，Resnet[8]模型的规模和深度也越来越大，因此机器学习领域从神级网络算法中诞生了一个独立的学派：深度学习，并且发展迅速，以至于现在人们谈论人工智能往往等同于深度学习。但是因为GPU硬件发展迅速，显存大小同步增大，使得这类模型仍然可以放在单块显卡下，我们可以通过 ZeRO optimizer[25]，CPU-offload 等技术进一步完善 GPU 使用效率显著，甚至有文章[26]说实用模型总能在单卡上放下并训练。但即便如此，模型的参数仍然有数亿的规模，需要大量的数据来训练。为了加快训练，我们仍然需要使用它数据并行对于分布式训练，由于感知模型通常是密集的参数，几乎所有的参数都需要更新，因此通信行为与稀疏模型非常不同，所以我们有另一种并行分布式：AllReduce同步数据并行[27]。事实上，这种分布式训练模式来自高性能集群 HPC 在该领域，即每个培训参与者维护一个模型的副本，然后在培训过程中将每个工作节点分配到 N 通过前向和反向计算，得到该数据下模型估计值与目标指数之间的差距，从而得到该数据上参数需要变换的梯度，然后使用高性能 HPC 集群通讯库原语 AllReduce 使每个节点同时加权所有节点的局部梯度变化，使每个工作节点得到全局梯度变化的副本，最后将梯度变化应用于模型参数，使每个工作节点上的模型副本同步更新。AllReduce 可细分为 Ring-AllReduce[26]和 Reduce-Scatter 其中 Ring-Allreduce 由于消息传递的分散化和非常确定性， Nvidia 可以根据这个特征构建多个特征 GPU 物理拓扑卡NVLink[9]，从而高效进行 Ring 的 Allreduce，并将这个通信原语集成到自己的软件中 NCCL 库[10]大大促进了数据并行模型训练的同步发展和人工智能应用的发展。当然，还有其他研究，如字节跳动 BytePS[11] 参数服务器的分布式架构用于密集模型 AllReduce 这个计算不好 NVLink 拓扑还可以进行高效同步数据并行训练。

随着深度学习的演变，谷歌在搜索推荐广告领域也从传统的机器学习算法演变为深度学习，如原始逻辑回归。谷歌在这一领域培训了最具影响力和创建深度学习模型的模型 2016 年发表于 DLRS 的《Wide & Deep Learning for Recommender System》[12]。为了有效支持感知和推荐模型的需求，谷歌团队推出了 TensorFlow 深度学习框架[13]可以统一多种并行分布式训练范式 Tensor 计算引擎上的流图深度学习。

虽然 TensorFlow 从一开始就对整个分布式深度学习训练进行了非常基本和基本的抽象，可以支持任何大规模的分布式训练方法，并在此框架的基础上迅速 2018 年就出现了流水并行[14]和算子拆分[15]模型并行。平行流水将模型层和层横向分为多个阶段，从而将模型的不同部分分配到不同的节点（GPU）因此，我们可以单独训练 GPU 无法存储的大模型。不同的模式 GPU 传输是层与层之间的计算结果，但我们知道整个神经网络的训练过程是前向和反向的计算过程，所以这些 GPU 因为计算依赖依次计算(图) 2 中的 b 所示)计算不能平行，不能充分利用 GPU 的算力。所以 Gpipe[14]进一步把计算进行 mini-batch 从而进一步让计算在切割中 GPU 类似微流水(图) 2 中 c 所示)并行进行。进一步提高 GPU 的利用率。另一种并行方法是算子分割，即纵向分割神经网络(如图所示 3 所示)，将算子计算切割到不同的节点进行计算，以克服计算依赖引起的流水 GPU 闲置问题，但不同的算子分布式计算需要不同的计算和拆分方法才能有效地并行化，并行效率也与算子类型密切相关，因此，这种方法往往需要大量的开发成本来平行训练模型。事实上，大规模推荐稀疏模型也是一种分割算子的方式，但这种算子是一个简单的稀疏特征 lookup 算子而已。事实上，大规模推荐稀疏模型也是一种分割算子的方式，但这种算子是一个简单的稀疏特征 lookup 只是算子。虽然研究人员很快就在那里 TensorFlow 在这个框架的基础上，给出了各种并行的方法来训练任何大型模型，但因为 GPU 硬件开发也很快，GPU 显存大小也发展迅速。因此，在很长一段时间内，除了推荐的大规模稀疏模型训练外，感知模型训练只需要数据并行训练，直到 2019 每年都有重量工作《BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding》[16]发表。

Bert促进超大规模预训练模型的快速发展 transformer 基模型结构为基础，将感知模型的参数规模从1亿级快速进入1000亿和10000亿级，如 GTP-2[17]，GTP-3[18]，T5[19]，M6[20]等等。最后，复杂的模型从学术研究领域并行进入企业实际实施过程。目前的超大规模模型往往需要根据模型的特点进行多种并行训练。这也是大家有时候叫的 3D，4D，5D 事实上，分布式训练本质上是将所有这些技能结合在一起。

因此，分布式并行训练的范式分为

• 大规模稀疏模型训练：参数服务器模式，异步训练
• 大规模稠密模型训练

1. 数据并行：AllReduce 或者参数服务器用极端显存优化梯度汇总（CPU-offload，Zero 梯度优化器优化，checkpoint 等)进一步扩大单节点的模型尺寸
2. 流水并行
3. 算子切分
4. 混合并行

大规模并行分布式训练的发展方向

自动分布式训练：模型混合并行越来越重要，但混合并行往往需要根据模型的特点和集群资源的现状，在模型的不同部分选择不同的并行模式。如果这需要模型构建者的选择，整个问题是一个问题 NP 为了更系统地解决这一混合并行寻优问题，问题将给模型构建者带来巨大的研发成本，使快速算法研发难以迭代，产业研究界进行积极相应的研究，如 GShard[21]，OneFlow[22]，Whale[23]等。其核心思想是使用静态图逻辑地描述网络训练步骤，然后进行系统的分布式训练，使用反映实际约束的成本分析模型（Cst Model）对于不同并行策略进行评估，最后自动选择最优的并行方案。

这些工作中，GShard 更多关注于谷歌的 TPU 集群，OneFlow 从重新构建一个面向分布式的深度学习生态，而 Whale 其实采取的理念是让算法同学利用现有深度学习生态，如同构建一个单卡的模型训练的方式去构建，并且能够在异构资源（不同能力的加速卡能力）环境找到更优的混合并行方案。图 4 展示了 Whale 的整体架构图。Whale 会把模型训练脚本转换成一个的 DAG 静态图，然后进一步转化成自有中间代码（Intermediate Representation，IR）描述的逻辑执行计划，然后系统进行自动规划，再根据规划出来的并行策略对执行计划进行切图，并且把系统资源进行相应的划分，再把切图的结果和资源绑定，最后让 Cost Model 在多个可能并行策略上进行选择，从而自动寻找到接近最优的分布式策略。

更大规模异构多模态多任务的模型训练: 随着模型算法，数据以及算力的发展，研究者终于可以开始构建神经网络算法最初的设想，就是建立类似大脑多任务综合在一起学习的方式，从而使得机器学习能够进一步仿真人脑的认知过程，把语言，视觉，听觉等多个认知任务关联在一起来进行整体训练，这个就要求更加异构的训练过程。Pathways[24] 就是在这个方向最新的探索。

总之，随着互联网的蓬勃发展，人们开始积累非常多的有价值的数据，同时随着半导体行业的方式使得我们具备足够算力使得模型研究者能够更敢去构建更加大型的模型，从而大大提升了模型的效果，产生能够媲美人类认知的人工智能模型，而大规模分布式训练的工程成为能够粘合数据，算力和算法的关键。

参考文献

[1] KunPeng: Parameter Server based Distributed Learning Systems and Its Applications in Alibaba and Ant Financial. KDD 2017.

[2] Persia: An Open, Hybrid System Scaling Deep Learning-based Recommenders up to 100 Trillion Parameters.CoRR 2021.

[3] Scaling Distributed Machine Learning with the Parameter Server. OSDI 2014.

[4] Parameter server for distributed machine learning. NIPS 2013.

[5] More effective distributed ml via a stale synchronous parallel parameter server. NIPS 2013.

[6] ImageNet Classification with Deep Convolutional Neural Networks. NIPS 2012.

[7] Going deeper with convolution. CVPR 2015.

[8] Deep Residual Learning for Image Recognition. CVPR 2016.

[9] https://www.nvidia.com/en-us/data-center/nvlink/

[10] https://developer.nvidia.com/nccl

[11] A Unified Architecture for Accelerating Distributed DNN Training in Heterogeneous GPU/CPU Clusters. OSDI 2020.

[12]Wide & Deep Learning for Recommender System. DLRS 2016.

[13] https://www.tensorflow.org/        

[14] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. NIPS 2019.

[15] Mesh-TensorFlow: Deep Learning for Supercomputers. NIPS 2018.

[16] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL-HLT 2019

[17] https://huggingface.co/docs/transformers/model_doc/gpt2

[18] https://openai.com/blog/gpt-3-apps/

[19] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. Mach. Learn. Res. 2021.

[20] M6: Multi-Modality-to-Multi-Modality Multitask Mega-transformer for Unified Pretraining. KDD 2021

[21] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. CoRR 2020.

[22] OneFlow: Redesign the Distributed Deep Learning Framework from Scratch. CoRR 2021

[23] Whale: Scaling Deep Learning Model Training to the Trillions. https://arxiv.org/pdf/2011.09208.pdf

[24] https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/

[25] ZeRO: memory optimizations toward training trillion parameter models. SC 2020

[25] DeepSpeed: https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/

[26] https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/

[27] Horovod: fast and easy distributed deep learning in TensorFlow. arXiv preprint arXiv:1802.05799 (2018)