要了解大模型训练难,我们得先看看从传统的分布式训练,到大模型的出现,需要大规模分布式训练的原因。接着第二点去了解下大规模训练的挑战。
从分布式训练到大规模训练
常见的训练方式是单机单卡,也就是一台服务器配置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能够在大模型引领业界。