从分布式训练到大模型训练

要了解大模型训练难，我们得先看看从传统的分布式训练，到大模型的出现，需要大规模分布式训练的原因。接着第二点去了解下大规模训练的挑战。

---

从分布式训练到大规模训练

常见的训练方式是单机单卡，也就是一台服务器配置1块AI芯片，这是最简单的训练方式。随着数据量的增加，希望加快模型的训练速度，于是出现了单机多卡，多块AI芯片并行，以一台机器上配置8块AI芯片为例，把数据切分成8份，分别在8块AI芯片上都跑一次BP算法，计算出梯度，然后所有AI芯片上计算出的梯度进行平均，更新模型参数。这样的话，以前一次BP只能训练1个batch的数据，现在就是8个batch。

模型和数据规模的增大，意味着训练时间的增长。为了提升模型训练的速度，可以增加计算资源来缩短训练时间，于是出现了分布式训练。所谓分布式训练，物理上就是使用多台机器，每台机器上都有多块AI芯片，网络模型运行在不同机器的不同AI芯片上，以加快整体训练速度。简单来理解，实质就是将单卡的负载拆到了多卡上。

来看看引入分布式训练的两个原因：

1）第一个是数据规模大，导致训练时间很长。在单机8卡情况下，对于COCO在115k规模的数据集，假设训练resnet152模型需要40个小时。训练Open Image dataset v4图片数据集在1,740k规模下，则相对来说可能需要40天时间。对于生命宝贵、毕业延期的炼丹同学来说，需要不断的调整网络结构、尝试新的参数，修改下游任务等工作，单机单卡和单机多卡的训练速度是无法接受的，于是很有必要使用分布式进行训练。

2）第二点就是分布式训练可能带来精度上的提升。先回忆一下，为什么要用SGD来优化模型，随机梯度下降的“随机”是指每次从数据集里面随机抽取一个小batch数据来进行计算真实值和模型预测值之间的误差，然后反向传播。之所以只选一个小batch，一是因为随机数据产生的小batch梯度方向，理论和实践证明，一定程度上可以模拟整个数据的梯度下降方向，二是因为AI芯片的内存有限。但是在实际情况下，小batch梯度并不足够代替整个数据集的梯度方向，每次小batch在BP算法求解出来的梯度方向，与整体数据集并不完全一致。这样就会导致优化迭代（训练）过程不断震荡。使用分布式训练，可以使用更大的batch size，避免优化过程中的震荡。

数据并行：在分布式训练场景，因为数据规模增大，可以通过数据并行，可以对输入数据进行切分，每块AI芯片只需要处理一部分数据。优点是计算通信可以重叠，通信需求低，并行效率高；缺点是无法运行大模型。

模型并行：同时，模型每一层的参数量增大、模型的层数越深，可以通过模型并行修改网络层内的计算方式，将单层的计算负载和内存负载切分到多块AI芯片上。缺点是通信和计算是串行的，对通信的要求高。

流水线并行：另外还可以利用流水并行，将不同的网络层放在不同的AI芯片上运行，进而进一步将计算负载和内存负载切分至多块AI芯片。优点是计算的通信可以重叠，通信需求低；但是存在流水线间空闲bubble，需要和重计算配合使用。

通过对数据、网络模型、以及运行的流水线进行切分，分布式训练可以有效地减少单卡的资源负载，一方面提升了训练任务的计算和内存吞吐量，另一方面使得原来单机单卡无法训练的任务，变成可能。

现在一切看上去都是那么的美好，有分布式训练就够了。理论上，AI芯片数量越多，模型训练越快。但是，随着训练数据集规模的进一步增长，数据并行就会出现局限性：当训练资源扩大到一定规模时，由于通信瓶颈的存在，增加计算资源的边际效应，就会越来越明显，甚至增加资源也没办法进行加速，这个比例称为加速比。假设单设备的吞吐量为T，有n个设备系统的吞吐量应为nT，设这个系统实际达到的吞吐量为T_n。最终的加速比(scale factor)为：

理论上，我们希望无论有设备数n为多少，加速比越接近1越好。

另外，随着网络模型的不断增大，会出现模型需要的内存急剧膨胀，算子的增加也会使得单卡就算进行模型切分和流水线切分，也难以在合理的时间内完成一个step的训练，最终导致分布式训练不再适用于这类型的任务。这也是大模型的出现，对系统的挑战。

这个时候，面对大模型，我们需要引入大规模训练技术，在解决内存的同时，也不会被计算资源给限制住，让算法开发人员可以方便进行高效的分布式调试调优和编写代码。愉快地训练大模型。

---

大规模训练的挑战

相比普通的分布式训练，大规模训练在技术上，需要考虑的问题更加复杂。

首先，面对单卡无法装载的大模型，如何利用多卡来突破内存限制的瓶颈是个问题；其次，大规模训练会用到大量的计算资源，大量计算资源间如何通信、协作是另一个难题；最后，如何平衡各类层出不穷的大规模训练技术，使得众多技术形成一个完整高效的训练方案，更是一大学问。

下面，将大规模训练技术面临的挑战分为四个部分：内存、通讯、计算和调优。

内存墙

模型训练无可避免的问题就是内存墙。

下面举两个例子，在进行ResNet152对120G的ImageNet数据进行训练时，在使用Batch Size=128，Input Image Size=512x512的前提下，训练过程占用了16G的内存，在GPU一般具有16G内存的情况下，没办法把模型放在一个一块AI芯片上。想进一步加大每次执行数据的Batch size，模型训练的内存占用也会随之增长，最后高于AI芯片的内存容量，触碰到了内存墙，模型无法训练。

以2000亿参数的鹏程盘古大模型为例，2000亿的参数内存占用就消耗了754GB，训练过程因为会有权重、激活、优化器状态、再加上自动微分所产生临时变量，需要3500GB的内存，一个大模型训练就需要100多块具有32G内存的AI芯片。

要理解内存墙出现的本质，就需要了解内存增长的本质。在ResNet50的1轮迭代，打开看看内存占用变化：

1）网络模型开始计算的时候，内存占用不断增加，直到达到峰值1.2GB。

2）峰值过后内存开始逐渐释放，内存占用慢慢降到320M。

3）1个step计算结束后，仍有一部分内存驻留，内存保持在320M。

可以发现，模型训练对内存的占用可以分为两部分：

1）一部分是模型自身的权重参数、优化器状态信息，由于是比较固定的所以称为静态参数。

2）另一部分是模型在前向计算和反向传播的时候，会产生诸如前向输出Forward Output、梯度输出Output Gradient、算子计算时候的临时变量等，这部分内存会在反向传播时逐渐释放掉，因此被称为动态参数。

从这里面可以发现，在计算过程的内存峰值，是否遇到内存墙，主要是由动态内存决定。即使绞尽脑汁，降低静态内存意义并不是很大，关键在于降低内存占用的峰值。

**静态内存**

以ResNet50为例，静态内存占用大约1.2G：

ResNet50 = 250M(权重参数) + 900M(优化器参数和动量) ~ 1.2G

在计算过程中，神经网络模型每一层的卷积或者全连接计算，都会把权重W_m长期保存下来，用作网络的权重参数更新。另外针对诸如ADAM的优化器，会存储优化器的动量等信息，用于优化器计算。

一块AI芯片有16G，最大能塞满20+亿参数的模型，但是这时候已经没有额外空间，留给动态内存进行分配啦。所以说降低动态内存的峰值才是人间正道。

**动态内存**

静态内存比较好理解，下面主要看看动态内存。动态内存一般指的在自动微分Autogard过程，产生的中间变量。

神经网络层数非常多，在计算过程中，神经网络模型每一层的前向输出Forward Output都需要保存下来给反向传播时候使用；而在神经网络的反向传播计算过程中，需要记录下权重梯度Weight Gradient(W_g)，给优化器使用，在后续的step中更新梯度的信息；另外还有梯度输出Output Gradient同样在反向传播的过程中，继续在反向传播的过程中，作为上一层网络模型的输入。

当然，这里面还会有一些正向和反向的算子中间变量需要消耗内存的，由此，上述的O_f、O_g、W_g、OP累积下来,会造成巨大的内存开销。从ResNet50的例子中，我们可以看出，动态内存是静态内存的X倍(，静态内存占用了1.2G内存，而动态内存则占用了25G内存，比例非常不平衡)。

静态内存和动态内存都可能造成内存墙的问题。相互独立但又相互制约，任意一侧的增大都会导致留给另一侧的显存空间变小，所以单单对一侧做优化是不够的，必须同时优化静态和动态内存。

为了能够将大模型运行起来，需要使用模型并行、流水并行等技术，但是这些技术又会降低AI芯片的计算吞吐量。

通讯墙

大模型通过模型并行、流水线并行切分到AI集群后，通讯便成了主要的性能瓶颈。

从模型切分的角度来考虑，大模型再怎么切分其仍然是一个模型，因而模型切分到不同的机器后，仍然需要通信来对分布在不同机器的参数进行总体聚合。

在通讯方式的角度来看，参数聚合(All Reduce)就会对通讯提出很多需求，例如使用同步的更新策略，还是使用异步的更新策略，如何对模型局部变量进行更新等。

另外，在网络通讯的角度，由于专用的AI加速芯片中内存与计算单元ALU非常进，使得片内带宽很大，计算速度很快，但是在集群中的网络传输速度远远不能匹配专用的AI加速芯片运算速率。

直接用更大的带宽不能解决这些问题吗？

从图中，我们可以看到，随着网络带宽从1Gbps到100Gbps，利用率从接近100%下降到40%左右。随着网络带宽的增加，带宽的利用率将会越来越低，高带宽所带来的收益会遇到瓶颈。

随着机器规模的扩大，基于同步的All Reduce通讯聚合方式，会因为大量的AI芯片和服务器之间频繁进行同步，出现水桶效应，也就是最慢的一路通讯，将会决定整个AI集群的通讯的高度。

如果采用目前比较流行的Ring-AllReduce的通信聚合方式，当通讯的环越大，通讯的延长将会不断地被扩大。另外网络协议的多次握手的方式，诸如此类的开销会导致训练无法有效利用带宽。

总体而言，通讯过程中，需要综合考虑数据参数量、计算量、计算类型、数据样本量、集群带宽拓扑和通讯策略等不同的因素，才能设计出性能较优的切分策略，最大化利用通讯效率，提高通讯比。

性能墙

性能墙呢主要是指计算资源利用率的问题。随着大模型的提出，对算力需求更加迫切，理论上在4K的集群上每块卡快1分钟，总体就快了68个小时。大模型会增加对算力的需求，但是随着大模型引入各项分布式并行技术的同时，会降低计算资源的利用率。

计算效率在MMlab的贴总结得非常好，我们同样从底至上分为算子层(Operator Level)、图层(Graph Level)和任务层(Task Level)三个方面。

**算子层(Operator Level)**

在Operator Level的问题大模型和普通模型遇到的问题都是类似，主要是指对算子级别的优化。

1）其中比较突出的问题是模型小算子过多，可以通过算子融合进行优化，例如对Conv+Bn+Act三个算子进行融合成一个大算子。但是还有很多算子融合策略无法穷举，如何通过计算编译技术，对小算子进行融合。

2）第二点就是算子实现不够高效，类似于卷积CONV算子针对2x2和3x3 Kernel大小的使用Winograd算法代替。可是这种算法是人为提出了，更多的高效Kernel算子，如何利用编译技术，去寻找最好的多线程并行计算。

3）最后一点就是内存局部性差，对算子和内存的开销进行分析，可以对计算时算子输出有相同的shape进行复用。

**图层(Graph Level)**

Graph Level指如何对计算图进行优化，进而加速大规模训练。

一般而言，如果没有遇到控制流和大模型，一个模型只有一个图。但是当遇到大模型的时候，由于一块AI芯片的内存没办法满足大模型的内存需求，因此需要对算子和模型进行切分。这时候涉及两方面的问题。

1）如何搜索和切分处计算效率更高的子图，分配到不同的机器上进行计算。

2）数据在通讯和内存之间如何增加overlay重叠部分，提高单卡整体的计算率。

**任务层Task Level**

大集群大模型的分布式混合了各种并行训练策略，在任务层级，性能的瓶颈从计算转移到了通信，如何降低通信量，缩减通信域规模，尽可能把通信时延隐藏在计算中，是大规模训练的核心关注点。

大规模训练技术中，不仅要求AI芯片的计算性能足够强悍，同时也依赖于AI框架的大规模分布式训练的运行和调度效率，以及在分布式并行等各种优化手段的权衡。

调优墙

在数千节点的集群上进行模型开发，听到就头皮发麻。我平时把EffectNet网络模型魔改，还不一定能够收敛，调一次参数再训练一次，单机多卡一个星期就过去了。那要是大模型训练了1个月遇到Loss跑飞了怎么办？

所以在数千节点的集群上，需要考虑到提升算法工程师分布式调试调优的效率，另外还要考虑降低工程师对大模型进行并行切分的难度。除了对人的考虑，还要对硬件集群的管理，需要保证计算的正确性、性能、可用性。要是有一台机器坏了，如何快速恢复训练中的参数。

---

总结一句话，大模型训练考验的是算法、数据、框架、资源调度等全栈和全流程的综合能力。希望MindSpore能够在大模型引领业界。

引用

1. Kolesnikov, Alexander, et al. "Big transfer (bit): General visual representation learning."Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part V 16. Springer International Publishing, 2020.
2. Gao, Yanjie, et al. "Estimating gpu memory consumption of deep learning models."Proceedings of the 28th ACM Joint Meeting on European Software Engineering Conference and Symposium on the Foundations of Software Engineering. 2020
3. Zhang, Zhen, et al. "Is network the bottleneck of distributed training?."Proceedings of the Workshop on Network Meets AI & ML. 2020
4. Agarwal, Saurabh, et al. "On the Utility of Gradient Compression in Distributed Training Systems."arXiv preprint arXiv:2103.00543(2021)
5. OpenMMLab：大规模训练系列之技术挑战