近几年,随着“大模型”概念的提出,深度学习模型越来越大,如何训练这些大模型成为一个亟待解决的工程问题。最初的视觉模型只有几百兆的参数量,而现在的语言模型中,动则百亿,千亿的参数量,甚至万亿的大模型也是见怪不怪。如此巨大的参数量将会消耗巨大的存储空间。如下表所示为当前大模型的参数量(以Float32计算)以及对应的存储空间。
而当前最好的nvidia GPU显卡也只有40G的显存容量,显然将大模型塞进一张显卡是不现实的。本质上,所有大模型的训练,都使用了分布式的方式。当前分布式训练中,常用的有数据并行,模型并行和流水线并行,从计算效率上来说,数据并行要远远优于模型并行和流水线并线。但是数据并行对显存的占用是最高的,因为它需要将整个网络都运行在一张GPU上面。而在模型训练过程中,除了参数以外,还有很多地方需要占用存储空间,这就使得训练大模型时候的显存消耗进一步提升。因此你是否好奇,如何训练如此庞大的深度学习模型呢?
在探讨如何进行大规模训练之前,我们先来详细看看网络中的显存占用。通常在深度学习训练过程中,涉及到的显存占用包括:网络的参数,梯度,激活值,激活值的梯度,优化器的状态信息,如果使用了混合精度[6]训练,那么还有备份参数(master_weight)等。这里需要指出的是,激活值的梯度在古老的caffe框架中是没有做过优化的,其占用空间和激活值相同。但是在tensorflow,pytorch等框架中,已经做了很好的优化,因此激活值的梯度实际上并没有占据很大的显存空间。另外,可能很多朋友对前面几种类型的显存占用(网络的参数,梯度,激活值,激活值的梯度)比较清楚,但是对于优化器的状态信息以及混合精度[6]训练的备份参数(master_weight)不是很清楚,这里稍加说明。
在简单的SGD优化器中,更新参数使用如下公式:
但是通常我们不会直接使用SGD来更新参数,而会对梯度进行滑动平均后,再进行更新,即使用Moment SGD优化器,其计算公式如下:
由于计算需要用到,因此需要一直保存在显存中。就是优化器的状态信息,其大小和梯度一致,因此和参数大小一致。例如参数规模是100亿,那么优化器缓存信息也是100亿的规模。
在很多时候,我们也会使用ADAM优化器进行参数更新,而ADAM会用到梯度的一阶矩估计和二阶矩估计,公式如下:
同理,由于计算需要用到,计算需要用到,因此和两个变量也需要一直保存在显存中,他们大小也和梯度一致,因此也和参数一致。例如参数规模是100亿,那么和一共就是需要200亿的规模。
有时候我们为了提升效率,会使用混合精度[6]进行训练,而混合精度[6]训练为了抵消float16在参数更新时候的舍入误差,会额外保存一份FP32的参数用于参数更新,称作master-weights,因此会多出一份显存占用空间。
上面分析了深度学习网络的显存占用,而不同的架构中,各个部分的显存占比是不同。在CNN中,通常是激活值占据了大部分的显存空间。而MLP/Transformer等结构则是模型的参数与参数的梯度,优化器的状态信息占据了更大的比例。因此对于不同的网络结构,其优化策略是不相同的。对于CNN网络,通常优化的重点是激活值。而MLP/Transformer则更关注于网络参数相关的显存优化。因此对于两种不同类型的网络,分别需要用到激活值优化策略与参数优化策略。
前面分析过,cnn网络的显存消耗主要是激活值以及激活值的梯度。随着输入分辨率的提升,以及batch size的扩大,激活值以及激活值梯度的显存占用会呈现平方倍的增加。因此对于CNN的大模型训练,主要集中在对激活值的优化上。这里介绍一种以增加计算时间来降低显存空间占用的方式:“亚线性内存优化[5]”。首先来看一下常规深度学习的流程,如下图所示。前向计算后,保存所有的激活值,如图中的a1,a2,a3与a4。反向计算的时候,根据之前的激活值,计算每一层layer的梯度,包括激活值梯度与参数梯度。
而“亚线性内存优化[5]”的深度学习流程则与传统的深度学习流程有些差异,如下图所示。在前向计算中,为了降低显存消耗,会选择性的丢弃部分激活值,例如a2。反向计算的时候,按照常规的方法进行计算,当遇到激活值缺失的时候,例如计算layer3的参数梯度时,需要用到激活值a2,但是a2已经被丢弃,此时会暂停反向传播,重新进行一次最短路径的前向计算,根据a1计算出a2(如果a1也被丢弃,那么继续向前找)。然后再继续原来的反向传播。
由此可见,“亚线性内存优化[5]”彻彻底底的使用了时间换空间的策略。那么你是否好奇,为什么它会取一个这样的名字呢?其实理解起来也很简单,假设每一层的激活值大小相同,那么整个网络的激活值大小就和layer的数目成线性关系,layer的深度扩大几倍,激活值的占用就扩大几倍。而使用了“亚线性内存优化[5]”策略之后,显存的占用与layer的增加不再是线性关系,而是亚线性的,因此叫做“亚线性内存优化[5]”,实际上,如何选择丢弃的激活值,将会极大的影响最终的网络性能。被丢弃的激活值需要满足前向计算简单(降低重计算的耗时),激活值占用空间大的特点。例如BN层,其计算非常简单,重计算基本不耗时(相对卷积),因此比较适合丢弃。
除此之外,还可以使用低精度训练的方式。这样激活值以及激活值的梯度都使用更小的数据格式存储。能够极大的降低激活值的存储空间。例如在使用混合精度[6]训练的时候,所有激活值使用float16的格式,相对于原始的float32,显存占用直接缩小了一半。
上面介绍的激活值优化策略适合于CNN这样的结构。而当前的一些大模型,其显存占用主要集中在参数以及参数相关的显存占用上。例如参数的梯度,优化器的状态信息等。因此如何优化参数相关的显存占用对于大模型的训练显得更加重要。
微软开源的DeepSpeed训练框架中,使用了一种称为ZeRO[4]的显存优化技术,称为零冗余优化技术。本质上,它是一种数据并行的分布式训练策略,重点优化了数据并行中的显存占用问题。在ZeRO[4]数据并行中,每个GPU上虽然拥有完整的网络,但是每个GPU只保存一部分的参数,梯度和优化器状态信息,这样就就可以将参数,梯度,优化器状态信息平均分配到多个GPU上。这对于参数观规模较大的网络,显存的降低将是巨大的。但是由于分布式存储参数,也会导致通信的增加。
首先来简单回顾一下传统的数据并行流程,包括前向计算,反向计算,参数更新三个流程。假设一共有两张GPU参与训练,前向流程如下图所示,起始阶段,每张GPU初始化为相同的参数,并划分互斥的训练子集。每站GPU独立完成所有layer的前向计算。
反向流程如下图所示。起始阶段,最后一个激活值已经通过loss求导得到了激活值的梯度。然后每张GPU独立的进行所有layer的反向计算。
参数更新流程如下图所示。先对所有GPU的梯度进行规约操作(求平均值),然后每张GPU独立的更新参数。(由于初始参数值相同,梯度规约后也相同,因此最后每张GPU上更新后的参数也相同)
ZeRO[4]数据并行有多个级别,分别是os级别(只对优化器状态做优化),os+g级别(对优化器状态+梯度做优化),以及os+g+p级别(对优化器状态+梯度+参数都做优化)。我们直接分析优化程度最高的os+g+p流程。首先看前向计算流程,在起始阶段,每张GPU只保存W/GPU_NUM的参数。这里假设网络有2个layer,一共两张GPU参与并行,因此GPU1只保存layer1的参数w1,GPU2只保存layer2的参数w2。在layer1的前向计算之前,由于GPU2没有layer1的参数,因此需要做一次w1的参数分发。然后进行layer1的前向计算。同样在进行layer2的前向计算之前,需要将layer2的参数w2进行一次分发,再完成layer2的前向计算。整体示意图如下图所示。
再来分析ZeRO[4]数据并行的反向传播流程。同样的,当layer反向计算前,都需要对参数进行分发。然后再进行反向传播计算。完成反向传播之后,会有一个梯度搜集的过程,例如GPU2需要保存w2对应的梯度g2,因此所有其他GPU将g2梯度发送给GPU2。GPU2上面得到各个GPU的g2梯度后,做规约操作并保存,得到g2~。其他GPU将会删除w2,g2。然后重复该流程,直到所有layer都完成反向传播计算。示意图如下图所示。
最后来分析一下ZeRO[4]数据并行的参数更新流程。由于梯度的规约操作在反向传播的时候已经做了,因此ZeRO[4]数据并行可以直接更新优化器的状态信息,然后更新参数,示意图如下图所示。
从前面的ZeRO[4]数据并行流程可以看出,ZeRO[4]数据并行在os+g+p级别(对优化器状态+梯度+参数都做优化)优化时,会有两次参数的分发(前向计算一次,反向计算一次)和一次梯度的搜集。而传统的数据并行只需要做一次梯度的规约。所以ZeRO[4]数据并行的通信消耗将是传统数据并行的3倍?其实不然,传统数据并行中,虽然只需要做梯度的规约操作,但是由于每张GPU都需要得到规约后的梯度,因此使用的时all-reduce的通信原语。而ZeRO[4]数据并行中,虽然有3此数据传输,但是只需要一对多分发参数或者多对一的梯度搜集,使用的是broadcast和gather的通信原语。而broadcast和gather的通信消耗基本相当,约为all-reduce的一半,因此最终ZeRO[4]数据并行在os+g+p级别上的通信开始时原始数据并行的1.5倍,而不是3倍。当使用os+g级别的优化或者os级别的优化,通信消耗与原始数据并行相当。关于分布式数据并行中的通信原语,通信消耗我将另外编写文章分析。
最后我们来分析一下ZeRO[4]论文中最最最niubility的这副图。如何将传统数据并行中需要120G显存的模型变成了只需要1.9G。不看不知道,一看还真吓一跳!分析这张图片前需要知道一个前提,那就是全部基于混合精度[6]训练,并且采用adam的优化策略。
首先表示模型参数规模为7.5B,即75亿个参数量。由于所有训练均基于混合精度训练,因此参数和梯度都使用float16存储,一个参数占用两个字节,参数和梯度的显存占用都是的2倍。K表达的梯度状态信息以及混合精度master_weight的显存占用是的多少倍。由于梯度状态信息和混合精度master_weight必须使用float32来存储,即一个参数占用4个字节的存储空间,并且adam中有两个状态信息,分别是梯度的一阶矩估计和二阶矩估计,所以K = (2 + 1)*4 = 12倍的。Nd=64表示使用64张GPU进行zero数据并行训练。首先看第一行的Baseline,传统的数据并行,那么每张GPU的显存消耗就是:(2+2+K)\= 120G。接着看第二行,使用os级别的优化,那么参数和梯度的大小没变,优化器状态+master_weight被平均分配到了所有GPU上,因此每张GPU的显存消耗就是31.4G;同理分析第三行,使用os+g的级别的优化,由于梯度数据也被均分到了所有的GPU上面,因此每张GPU的显存消耗就是16.6G;最后分析使用os+g+p级别的优化,参数也被均分到所有GPU上面,因此最后每张GPU的显存消耗就是1.9G。从上面的分析中可以看出,在使用os+g+p级别的优化中,每张GPU的显存消耗就是传统数据并行的Nd分之1,大白话说就是用多少张GPU,显存消耗就能降低多少倍。
上面就是关于大模型训练训练中显存占用的一些优化措施。包括针对激活值优化的策略和针对参数优化的策略。正是有了这些工程上强有力的措施,才能让大模型的训练成为可能。不过由于作者水平有限,时间仓促,难免会有纰漏,还望各位读者不吝指正。感谢。
[1]DeepSpeed之ZeRO系列:将显存优化进行到底
https://zhuanlan.zhihu.com/p/513571706
[2]2022,大模型还能走多远
https://www.51cto.com/article/697186.html
[3]ZeRO+DeepSpeed:微软发布的高效大规模训练套件(含详细分布式训练流程)
https://zhuanlan.zhihu.com/p/108571246
[4] Rajbhandari S , Rasley J , Ruwase O , et al. ZeRO: Memory optimizations Toward Training Trillion Parameter Models[C]// SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. 2020.
[5] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training Deep Nets with Sublinear Memory Cost. arXiv preprint arXiv:1604.06174, 2016.
[6] Micikevicius, P. , Narang, S. , Alben, J. , Diamos, G. , Elsen, E. , & Garcia, D. , et al. (2017). Mixed precision training.