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

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

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

常见的训练方式是单机单卡，也就是一台服务器配置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)为：

$scale factor=T_n/nT$

理论上，我们希望无论有设备数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能够在大模型引领业界。