首页 A1训练模型 如何在普通显卡上训练“万亿大模型”?

如何在普通显卡上训练“万亿大模型”?

标题:如何在普通显卡上训练一个“万亿大模型”?

1.背景近年来,随着“大模型”概念的引入,深度学习模型变得越来越大。如何训练这些大型模型已成为亟待解决的工程问题。最初的视觉模型只有几百兆的参数,但现在语言模型拥有数百亿、数千亿甚至数万亿的参数并不罕见。如此庞大的参数量会消耗巨大的存储空间。下表显示了当前大模型的参数(以Float32计算)以及对应的存储空间。

目前最好的Nvidia GPU显卡只有40G显存容量。显然,将大型模型塞进一张显卡是不现实的。本质上,所有大型模型都是以分布式方式训练的。目前的分布式训练中,常用的有数据并行、模型并行和管道并行。就计算效率而言,数据并行远远优于模型并行和管道并行。然而,数据并行消耗的视频内存量最大,因为它需要整个网络在一个GPU 上运行。在模型训练过程中,除了参数之外,还有很多地方需要存储空间,这进一步增加了训练大型模型时的内存消耗。那么你是否好奇,如何训练这么大的深度学习模型呢?

2.深度学习中的视频内存使用。在讨论如何进行大规模训练之前,我们先详细了解一下网络中的显存使用情况。通常在深度学习训练过程中,涉及的内存使用包括:网络参数、梯度、激活值、激活值的梯度、优化器状态信息,如果使用混合精度[6]训练,还有备份参数。 (master_weight)等。这里需要指出的是,旧的caffe框架中并没有对激活值的梯度进行优化,它与激活值占用相同的空间。不过在tensorflow、pytorch等框架中,已经做了很好的优化,所以激活值的梯度实际上并没有占用大量的显存空间。另外,很多朋友可能对之前的显存使用类型(网络参数、梯度、激活值、激活值的梯度)比较清楚,但是对于优化器的状态信息以及混合精度的备份参数(master_weight)[ 6]训练)不是很清楚,这里稍微解释一下。

2.1 SGD优化器在简单的SGD优化器中,使用以下公式来更新参数:

2.2 Moment SGD优化器但通常我们不会直接使用SGD来更新参数,而是对梯度进行滑动平均,然后进行更新。即使用Moment SGD优化器,计算公式如下:

编辑切换到主页

由于需要进行计算,因此需要保存在显存中。它是优化器的状态信息。它的大小与梯度一致,因此与参数大小一致。例如,如果参数大小是100亿,那么优化器缓存信息也是100亿。

2.3 ADAM优化器很多情况下,我们也会使用ADAM优化器来更新参数,ADAM会使用梯度的一阶矩估计和二阶矩估计。公式如下:

编辑切换到中心

同样的,既然要使用计算,要使用计算,那么和两个变量也需要保存在显存中。它们的大小也与梯度一致,因此也与参数一致。例如,如果参数大小为100亿,则总大小为200亿。

2.4 混合精度[6] 有时为了提高效率,我们会使用混合精度[6]进行训练。为了抵消参数更新时float16的舍入误差,混合精度[6]训练会额外保存一份FP32参数。用于参数更新,称为master-weights,因此会占用较多显存空间。

编辑切换到中心

2.5 激活值优化和参数优化上面分析了深度学习网络的显存占用,在不同的架构中,各部分占用显存的比例是不同的。在CNN中,激活值通常占据大部分显存空间。 MLP/Transformer等结构是模型的参数和参数的梯度,优化器的状态信息占据较大比例。因此,对于不同的网络结构,优化策略是不同的。对于CNN 网络,通常优化的重点是激活值。 MLP/Transformer更关注与网络参数相关的内存优化。因此,对于两种不同类型的网络,需要分别采用激活值优化策略和参数优化策略。

3.激活值优化策略3.1 以时间换空间正如前面分析的,CNN网络的显存消耗主要是激活值和激活值的梯度。随着输入分辨率的提高和批量大小的扩大,激活值和激活值梯度的内存使用量将按平方因子增加。因此,大型CNN模型的训练主要集中在激活值的优化上。这里我们介绍一种通过增加计算时间来减少内存空间占用的方法:“次线性内存优化[5]”。首先我们看一下传统的深度学习流程,如下图所示。前向计算后,保存所有激活值,如图中的a1、a2、a3、a4。反向计算时,根据之前的激活值计算各层的梯度,包括激活值梯度和参数梯度。

编辑切换到中心

“次线性记忆优化[5]”的深度学习过程与传统的深度学习过程有些不同,如下图所示。在前向计算中,为了减少内存消耗,一些激活值,例如a2,会被选择性地丢弃。逆向计算时,按常规方法进行计算。当激活值缺失时,例如计算layer3的参数梯度时,需要使用激活值a2,但a2已被丢弃。这时,反向传播就会暂停。再次进行最短路径的正向计算,根据a1计算a2(如果a1也被丢弃,则继续向前查找)。然后继续原来的反向传播。

编辑切换到中心

可以看出“亚线性内存优化[5]”完全采用了时间换空间的策略。那么你好奇它为什么有这样的名字吗?其实很简单理解。假设每一层的激活值相同,那么整个网络的激活值与层数线性相关。层的深度增加数倍,激活值的占用扩大数倍。使用“亚线性内存优化[5]”策略后,显存占用与层数增加的关系不再是线性的,而是次线性的,因此被称为“亚线性内存优化[5]”。事实上,如何丢弃激活值的选择会极大地影响最终的网络性能。丢弃的激活值需要满足前向计算简单(减少重新计算耗时)和激活值占用空间大的特点。比如BN层的计算非常简单,重新计算基本不耗时(相对于卷积),所以更适合丢弃。

3.2 低精度训练此外,还可以采用低精度训练。这样,激活值和激活值的梯度就以更小的数据格式存储。可以大大减少激活值的存储空间。例如,当使用混合精度进行训练时[6],所有激活值都使用float16的格式。与原来的float32相比,内存占用直接减少了一半。

4.参数优化策略上面介绍的激活值优化策略适用于CNN等结构。对于目前的一些大型模型来说,其显存占用主要集中在参数以及参数相关的显存占用上。例如参数的梯度、优化器的状态信息等。因此,如何优化与参数相关的内存使用对于大型模型的训练来说更为重要。

4.1 ZeRO[4]数据并行原理微软开源的DeepSpeed训练框架采用了一种名为ZeRO[4]的显存优化技术,称为零冗余优化技术。本质上,它是一种数据并行分布式训练策略,重点优化数据并行中的内存使用问题。在ZeRO[4]数据并行中,虽然每个GPU都有完整的网络,但每个GPU只保存一部分参数、梯度和优化器状态信息,从而可以将参数、梯度和优化器状态信息平均分配给多个GPU。对于参数视图较大的网络,视频内存的减少将是巨大的。但由于分布式存储参数,也会导致通信量的增加。

4.2 传统数据并行流程首先简单回顾一下传统数据并行流程,包括正向计算、反向计算和参数更新。假设总共有两个GPU参与训练。前向流程如下图所示。在初始阶段,每个GPU都被初始化为相同的参数,并分为互斥的训练子集。每个站点的GPU独立完成所有层的前向计算。

编辑切换到中心

其逆过程如下图所示。初始阶段,通过loss对最后一个激活值进行求导,得到了激活值的梯度。然后每个GPU独立地执行所有层的逆向计算。

编辑切换到中心

参数更新流程如下图所示。首先对所有GPU的梯度进行归约操作(平均值),然后针对每个GPU独立更新参数。 (由于初始参数值相同,梯度下降后也相同,所以最终每个GPU上更新的参数也相同)

编辑切换到中心

4.3 ZeRO[4]数据并行流程ZeRO[4]数据并行有多个级别,分别是os级别(只优化优化器状态)、os+g级别(优化优化器状态+梯度)、os+g+p级别(优化优化器状态+梯度+参数)。我们直接分析最优化的os+g+p流程。首先看前向计算过程。初始阶段,每个GPU只保存W/GPU_NUM的参数。这里假设网络有两层,总共有两个GPU参与并行。因此,GPU1只保存了layer1的参数w1,GPU2只保存了layer2的参数w2。在进行layer1的前向计算之前,由于GPU2没有layer1的参数,因此需要对w1进行参数分配。然后进行layer1的前向计算。同样,在进行layer2的前向计算之前,需要对layer2的参数w2进行一次分布,然后完成layer2的前向计算。整体示意图如下图所示。

编辑切换到中心

我们来分析一下ZeRO[4]的数据并行反向传播过程。同样,在反向计算该层之前,需要对参数进行分布。然后进行反向传播计算。完成反向传播后,会有一个梯度收集过程。例如GPU2需要保存w2对应的梯度g2,因此所有其他GPU都将g2梯度发送给GPU2。获得GPU2上各个GPU的g2梯度后,进行约简操作并保存,得到g2~。其他GPU会删除w2、g2。然后重复这个过程,直到所有层都完成反向传播计算。示意图如下所示。

编辑切换到中心

最后我们来分析一下ZeRO[4]数据并行性的参数更新过程。由于反向传播时已经进行了梯度降低操作,ZeRO[4]数据并行性可以直接更新优化器的状态信息,进而更新参数。原理图如下图所示。

4.4 ZeRO[4]数据并行通信开销从前面的ZeRO[4]数据并行流程可以看出,ZeRO[4]数据并行在os+g+p层面进行了优化(优化优化器状态+梯度+参数)时,会有两种参数分布(一种正向计算,一种反向计算)和一种梯度集合。传统的数据并行只需要执行一次梯度下降。那么ZeRO[4]数据并行的通信消耗会是传统数据并行的3倍?事实上,在传统的数据并行中,虽然只需要梯度约简操作,但由于每个GPU都需要获取约简后的梯度,因此使用了all-reduce通信原语。在ZeRO[4]数据并行中,虽然有3个数据传输,但只需要一对多的分布参数或多对一的梯度收集,使用广播和聚集的通信原语。广播和聚集的通信消耗基本相同,大约是all-reduce的一半。因此,最终在os+g+p级别通信开始时的ZeRO[4]数据并行度是原始数据并行度的1.5倍,而不是3倍。当使用os+g级优化或者os级优化时,通信消耗相当于原来的数据并行度。关于分布式数据并行中的通信原语和通信消耗,我会另外写一篇文章来分析。

4.5 ZeRO[4]论文原图分析最后我们来分析ZeRO[4]论文中最nuble的图片。如何将传统数据并行需要120G显存的模型变成只需要1.9G。没看到你不知道,但是当你看到的时候真的很震惊!在分析这张图之前,需要知道一个前提,那就是都是基于混合精度[6]训练,采用adam的优化策略。

编辑切换到中心

首先,这意味着模型参数大小为7.5B,即75亿个参数。由于所有训练均基于混合精度训练,因此参数和梯度均使用float16 存储。一个参数占用两个字节,参数和渐变的显存占用是其两倍。 K表示的梯度状态信息和混合精度master_weight的内存占用是多少倍。由于梯度状态信息和混合精度master_weight必须使用float32存储,即一个参数占用4字节存储空间,而adam中有两个状态信息,分别是一阶矩估计和二阶矩梯度的估计,所以K=(2 + 1)*4=12 次。 Nd=64表示使用64个GPU进行零数据并行训练。首先看第一行的Baseline。采用传统的数据并行方式,每个GPU 的内存消耗为:然后看第二行,使用os级别的优化,那么参数和梯度大小都没有改变,优化器状态+master_weight均匀分配到所有GPU上,所以每个GPU的内存消耗是31.4G;同样分析第三个OK,使用os+g级别优化。由于梯度数据均匀分布到所有GPU上,因此每个GPU的内存消耗为16.6G。归根结底,采用了os+g+p级别的优化,参数也均匀的分配到所有GPU上,所以最终每个GPU的显存消耗为1.9G。从上面的分析可以看出,采用os+g+p级别优化时,每个GPU的显存消耗是传统数据并行的1/Nd。用白话来说,就是使用更多的GPU可以减少显存消耗。多少次了。

5.总结以上是大模型训练中显存使用的一些优化措施。包括针对激活值优化的策略和针对参数优化的策略。正是有了这些强有力的工程手段,才使得大型模型的训练成为可能。但由于作者水平有限,时间仓促,难免有错误,希望读者不吝指正。感激的。

6、参考[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 等人。 ZeRO: 针对训练万亿参数模型的内存优化[C]//SC20: 高性能计算、网络、存储和分析国际会议。 2020.

[5] 陈天琪,徐兵,张驰远,卡洛斯·格斯特林。使用次线性内存成本训练深度网络。 arXiv 预印本arXiv:1604.06174,2016。

[6] Micikevicius, P.Narang, S.Alben, J.Diamos, G.Elsen, E.Garcia, D. 等人。 (2017)。混合精准训练。

热门文章