万字长文解析:大模型需要怎样的硬件算力

万字长文解析:大模型需要怎样的硬件算力

首页冒险解谜Residual更新时间:2024-07-30

B站系统部服务器团队

负责服务器硬件引入与迭代,服务器资源交付、服务器硬件与系统运维。

郁晨

哔哩哔哩资深开发工程师

背景

2017年,《Attention Is All You Need》的一声惊雷,给世人带来了Transformer模型。2022年,以Transformer结构为核心的chatGPT,凭借其震惊世界的表现,掀起了“AI元年" 的序幕:从此,各方纷纷投入大模型领域,前赴后继,热情高涨。然而,随之而来地,大模型在软硬件基础设施、算法及数据集等多方面的挑战和困难也逐步显现;今天我们将从AI-Infra角度出发,带大家了解目前大模型在算力、存储、通信等多方面上硬件资源的挑战和瓶颈,并且提供相关指标的量化手段和选型指导。

二、LLMs引发的硬件瓶颈

在Transformer及大语言模型(LLMs)出现前,绝大部分的AI模型训练和推理,对算力、显存等硬件资源要求不高,使用单机CPU/GPU或分布式小集群即可满足需求。但LLMs的出现,在算力、显存和通信三个方向,都带来特别大的压力:

1) 算力瓶颈:

2)显存瓶颈:

3)通信瓶颈:

三、LLMs硬件需求量化

上述的分析,让大家对于LLMs的硬件需求有了一个简单认知。但关于每种规模的LLM究竟需要多少算力、显存和互联带宽等问题,可能还未特别明确。为了让大家产生具象而客观的认知,本节尝试量化了LLMs所需的大部分硬件资源指标。

3.1 模型参数量估算

3.1.1 为什么估算模型参数量?

准确地来说,模型参数量并不是一种硬件参数指标。但模型的参数量很重要,它和整个模型的算力、显存及通信带宽等主要硬件资源量化估算都紧密相关。因此,本小节先给出模型参数量的估算公式和推导过程,以便读者顺利理解后续的章节。

3.1.2 怎么估算模型参数量?

概括来说,大模型整体参数量约为 , 其中l为 tranformer块的个数,单个 transformer块的模型参数量约为 ,h 为Transformer块中隐藏层维度数。

3.1.2.1 分析推理

目前所有大模型,以Transformer模型为核心的,各种大模型的结构差异不大。总的来说,LLMs存在一定共性和规律,其结构如图1所示:一个大模型通常是由多个Transformer块串联而成,每个 Transformer 块中都包含固定的结构和计算顺序;

(图1. LLM结构示意图,source from[2])

基于此,LLM模型的参数量计算可以做到层层剥离:一个完整的模型共由l个transformer块串行而成。此外,在首块transformer之前和尾块transformer后还存在一定的结构。因此,我们可以将整个模型参数统计分为3部分:1.transformer块结构( )、2.首块transformer之前的结构(Vh)、 3.尾块transformer之后的结构(0)。当h较大时,第一部分参数量中的 的量级远高于剩余部分,所以整体可改写为 ;

接下来,我们就来具体分析下三部分参数量是如何推导的(参考[6]):

一、 串联的Transformer块结构共有 的参数量:

整个模型共存在l个transformer块,每个 transformer 块中包含的 的参数量。而对于每个transformer块而言,其参数计算过程如下:

1)首先,多头自注意力结构共包含 参数量:1)该结构一共有 a 个头。每一个头都拥有一个self-attention结构,其中包含 、、 三个矩阵,用于生成 Q、K、V 三个矩阵。而 、、 每个矩阵形状是(h,h/a)。因此: 、、 三种矩阵的参数量是 。此外每种矩阵还包含了h个偏置。总计;2)另外,后面有个 WO矩阵形状是(h,h),参数量是 ,还包含了h个偏置。总计 。综上,多头自注意力的结构中共包含 参数;

(图 2. Transformer 中多头自注意力计算流程示意图,source from[3])

2)其次, MLP中共包含 的参数:MLP有两个矩阵和一个激活函数。第一个矩阵形状是(h,4h),另一个矩阵形状是(4h,h),合计 的参数量。此外,两个矩阵中还分别包含了 4h 和 h 的偏置;

3)最后,layer norm中包含4h的参数:包含了2个可训练模型参数:缩放参数λ和平移参数 β;而在Self-attention 和MLP两个结构中各有一个layer norm:每个layer norm 中都是对于 b*s个tokens 进行 h 个数的归一化,所以有h个参数对(λ和 β)。故一个transformer块中共需要 个参数;

二、 首块transformer之前的结构的参数量为 :

在首个tranformer block之是input embedding的结构,其中包含了word embedding、 postional encoding两个结构:

1)Word embedding主要包含一个(V,h)的矩阵。用于将tokens 从词典的独立热编码维度(one hot vector)转换到隐藏层维度上。所以Word embedding结构中有 的参数;

2)Postional encoding 的参数可以忽略不记。这是因为具体两种情况参数量都非常少:1)如果不可训练(例如使用正弦、余弦函数),则不存在参数;2)如果可训练,参数量很少可以忽略;

三、 尾块transformer之后的结构的参数量约为0:

在最后一个tranformer block 结尾还可能还衔接了一个线性层矩阵和Softmax结构:1)线性层,对应了一个(h,V)的矩阵,用于生成的 token 数据从隐藏层维度转回到词典维度。这里的线性层就是之前 word embedding矩阵的转置,所以不涉及参数和训练;2)Softmax结构也基本无参数量;

3.1.2.2 结论验证

论文[4]给出类似的结论,模型参数量P计算公式为:

其中s 表示 tokens序列长度,h为隐藏层维度数,V代表词典维度数。对比论文[4]中公式和我们分析出的公式,发现两者基本一致:只相差了sh的参数量,sh量相比billion级别的模型参数量而言是微小的。

再尝试将一些主流模型代入试算公式 ,发现是公式和实际结果基本保持一致:(其中llama系列的模型结构和标准Transformer结构差异不小,但依然在可接受的误差范围内)

(表1. LLMs模型参数真实参数量和估算结果的比较)

注:Llama2-34B、Llama2-70B,非早期标准transformer-decoder结构。参数量估算公式为

3.2 模型算力需求估算

3.2.1 为什么要估算模型算力需求?

LLMs的结构呈现特定规律,已被证实能够相当准确地预测模型训练和推理所需的计算资源。因此,通过估算模型算力需求,我们能够有针对性地指导实际工作的开展:不仅有助于选择更加合适的硬件设备,还能评估出所需的算力卡数以及潜在硬件瓶颈等关键信息。这些相关内容,将在后续章节展开。

3.2.2 怎么估算模型算力需求?

模型的参数量和模型算力消耗之间存在一个简单换算关系:每个token输入集群时,单个参数量的训练需要6~8次浮点数运算,推理需要2次浮点数运算。接下来,我们看下这个结论是怎么一步步推导出来的。

3.2.2.1 通用矩阵乘法中的浮点计算次数

首先,我们可以发现,矩阵乘法(General Matrix Multiple)存在浮点计算次数规律的:如[m,k]的矩阵和[k,n]的矩阵相乘,需要 次浮点运算。公式中“2”是系数,表示矩阵乘法中的每个参数都需要约2次的浮点运算:1次加法和1次乘法。

而神经网络中的绝大部分运算都可以想象成是矩阵乘法,也可以应用上述给出的浮点计算公式。这里举个例子帮助理解:如图3。该图中输入层(x0~2)是3个神经元,输出层是2个神经元(y0~1)。那么前向传播过程可以看作一个[1,3]的矩阵和[3,2]的矩阵相乘,得到[1,2]的矩阵结果。图中灰色椭圆的"b0"、"b1"是偏置(Bias)。

(图 3. 一个简单的神经网络示意图)

在这个过程中,完成的运算如下:

上述2个公式中,每个公式都有3个加法、3个乘法计算。所以[1,3]的矩阵和[3,2]的矩阵相乘共包含 次浮点计算量(称为FLOPs)。从而推广到一般:[m,k]的矩阵和[k,n]的矩阵相乘,需要 FLOPs。

而GPU在不同精度单位下,浮点数运算的效率是不同的。例如,A100在16位精度(FP16或BF16)下每秒可以进行 次浮点运算。而在32位精度下,每秒只能进行 次浮点运算。所以,我们可以理所当然的想到一个提速思路 - 让所有运算在低精度下进行来提高训练/推理效率。但是,这里有个无法忽略的问题:某些计算必须要在高精度下进行。例如全程在低精度下进行模型训练,往往会出现由于精度不够导致数值溢出和无法收敛的问题。所以目前大模型训练通常采用的是混合精度方式[10]进行:如图4,前向传播、反向传播这些无需高精度的计算采用16位精度(FP16或BF16)进行,优化器状态更新在32位精度(FP32)进行。而模型推理的要求就比较低,可以在16位、8位甚至4位精度进行,整个任务的计算效率非常高。

(图4. 混合精度训练示意图,source from [10])

3.2.2.2 分析推理

先观察模型训练和推理的算力主要用在哪里了:

我们令模型需要浮点计算次数(FLOPs)为 ,而 前、反 是前后向传播分别需要的FLOPs。那么:

训练:前反前未开启激活值重计算时前反前开启全激活值重计算时推理:前

此时,我们只需要知道 前 ,就可以估算出训练和推理各自需要多少FLOPs。

接下来进行 前 的计算推导:

先确定符号:我们令B对应batch size,s对应sequence length,h对应hidden dimension,l对应layers number(更多信息请参考表3);接着,我们再来观察大模型结构:大模型本身由l个transformer块串联而成。此外,首transformer块前和尾transformer块后还有一些结构。所以大模型的算力估算可以分为两大块:一、l个transformer块中的FLOPs;二、其他结构中的FLOPs:

一、 l个transformer块中,每个transformer块都由多头注意力结构和MLP两种结构构成:

1)注意力结构中,

单个attention共包含 FLOPs。

2)MLP中,包含了[h, 4h]的矩阵,以及[4h, h]的矩阵,还有一个激活层。

单个MLP共有包含 FLOPs。

因此,每个transformer块算力需求为 。由于有l个transformer块,共计 FLOPs。

二、 其他结构中的FLOPs:

主要是输出层logit 层的计算。最后一个transformer block输出的激活值为[B,s,h],和[h,V]的线性层进行矩阵乘法:需要 FLOPs。

因此,前向传播浮点计算次数如下:

而通常在大模型中,6h远大于s,12lh远大于V,所以还可以进一步简化:

综上,可得单个LLM模型副本的训练和推理的算力需求如下:

训练:前未开启激活值重计算时前开启全激活值重计算时推理:前

结合3.1节对于模型参数量的评估公式 ,发现

1)模型训练中:如果没有应用激活值重计算,单模型副本处理每个tokens时,单参数上算力需求约为6FLOPs。如果应用全激活值重计算,提升至8FLOPs;

2)模型推理中:单模型副本处理每个tokens时,单参数上算力需求约为6FLOPs;

3.2.2.3 结论验证

英伟达&微软在[4]给出训练所需的算力:

与我们上述分析及公式基本一致。侧面反应了上述公式准确性。

我们尝试本节推出公式和LLMs模型的实际计算量进行对比,发现保持基本一致:

(表2. LLMs实际算力需求和公式估算结果的比较)

3.3 模型通信资源量估算

3.3.1 为什么要估算模型通信资源量?

LLMs运行中产生通信的原因,在第二节LLM引发的通信瓶颈中有说明,这里就不赘述了。

量化和归纳集群中的通信量、通信频率以及流量方向的等规则,让我们了解整个LLM集群内的机器内(Intra)和机器间(Inter)的通信压力(前者关联的是机内高速互联带宽,如NVLink、PCIe等。后者关联的目前主要是RDMA网络和网卡),从而更好地指导设备选型、组网设计和部署工作的开展,后续也更容易对自研通信的算法、网络流量进行针对性优化。

3.3.2 怎么评估模型通信需求?

无论是训练还是推理,当前LLM集群中通信规则基本确定。我们可以根据模型规模、并行度配置以及其他超参数来推算出集群中的流量:

1)目前主流模型训练是数据并行(Data Parallel)、张量并行(Tensor Parallel)、流水线并行(Pipeline Parallel)构成的3D并行。而在模型推理方向中TP、PP也是标配。3D并行中的通信规则是确定的;

2)各自做显存优化、算力优化的各种技术(如zeRO系列、序列并行[2]等),在用通信资源换取其他硬件资源时,引入的通信量和规则也是确定的;

3)集群中的集合通信被AI框架中各种并行技术、显存优化技术所调用,限定在特定几个原语(primitives)的范围之内,其底层的流量规则也是确定的。后续会写一个LLM集合通信相关的文章,所以这里不做过多展开;

在进行集群中LLM流量规则分析和流量量化前,我们先进行必要的参数说明。本文后续章节都会用得到:

(表3. 参数含义表)

3.3.2.1 并行技术

并行技术分为数据并行(Data Parallelism)和模型并行(Model Parallelism)。

一、 数据并行

1)简单数据并行

DP的示意图如图5所示:其核心思路很简单,将模型结构看成一个整体。通过复制多个模型副本,并将每个模型副本部署到不同设备上运行。

对于AI场景而言,DP仅在模型训练场景中使用。而推理中由于只需一份模型副本即可实现推理,所以DP不适用。

(图5. DP流程示意图,source from Nvidia)

数据并行工作流程是,在训练时每个迭代开始时,每个模型副本喂入不同的数据样本。在前向传播和反向传播都完成后,各个模型副本间以allreduce的通信原语进行梯度数据同步;而这里的allreduce同步其实有同步(Synchronous)和异步(Asynchronous)两种模式:

数据并行的优点和缺点都比较明显:

所以,当前LLM训练场景中,通常在保证模型能够塞下(通过TP和PP)的情况下,保证足够的数据并行度;

二、模型并行

模型并行指的是在每个模型副本内,将模型结构进行切分:模型的各个部分使用不同的算力节点来负责计算。由于是对具体的模型结构的拆分,而每种算法的模型结构各不不同,所以模型并行比简单数据并行要复杂。

而在LLM领域,由于是基于transformer结构的,具备共性。所以针对transformer结构的模型并行技术得以发展:

如图6所示,TP和PP都是大模型时代的模型并行的产物,对于大模型训练和推理都至关重要,解决了模型副本高效切割等关键问题。

(图6. 流水线并行和张量并行的示意图, source from OpenAI)

1)流水线并行

PP中,l个transformer块被切割成了p段,部署到p个GPU上,每个GPU称之为一个PP stage。每个stage大约有 个transformer块。如果仅仅引入上述的切分的思想,整个模型只会有一个batch输入:会发现整个系统非常高的硬件空闲率(如图7.b所展示的)。每个时刻都只会有1个算力卡在工作,剩下的算力卡都在闲置。通常,我们将硬件空闲率(idle rate)称之为气泡(bubble),如何降低气泡大小是各种PP技术需要解决的主要难题之一。

引入pipeline的概念可以提高硬件利用率(如图7.c所示):在一个物理模型副本上,将每次迭代所用的batch size B拆分成m个microbatch;此时,系统中可认为独立存在m个逻辑模型副本: 每个副本被喂入不同的训练样本(micro-batch),计算出不同的梯度数据。并在m个副本都完成前后向传播的任务后,就会进行逻辑副本间的梯度同步和参数更新,这就是Google的PP技术(Gpipe[7])的逻辑。而微软的PP技术(Pipedream[13])编排的更加复杂些,但也能获得更小的气泡。但两者其基本思想类似,所以就不展开讨论了。

因此,可以将流水线并行看作是数据并行和模型并行的混合产物。

(图7. Gpipe中流水线并行图示,source from [7])

我们再来看下PP内部的通信情况:

PP产生的通信发生在前后向传播内,

因此,一次迭代内两个相邻的stage GPU间发生了2bsh个参数量的send/recv。而一般PP编排的比较精妙,是可以将上述通信过程overlap到计算过程中(如图8所示)。我们这里不展开,在后续集群拥塞章节进行专项展开。

(图8. PP中的通信和计算的示意图,source from [13])

2)张量并行

TP是针对transformer结构的定制化拆解技术,利用的是1.transformer块多头注意力结构中每个头之间的独立性、2.矩阵乘法的可拆解性。

对每个transformer块中的attention结构和MLP结构都可以进行拆解:

(图9. 张量并行中自注意力结构的拆分,source from [8])

MLP是由两个线性层和一个激活函数构成的。两个线性层,可以做是两个矩阵:一个是[h,4h],另一个是[4h,h];可以将第一个矩阵按列拆分,第二个按行拆分。最后在做allreduce;类似地,这个过程中,前向传播和反向传播共涉及2次allreduce;

(图10. 张量并行中MLP结构的拆分,source from [8])

因此,从张量并行本身来看,每个transformer块需要进行4次关于激活值和激活值偏导数的allreduce,每次allreduce的通信量和bsh有关的。如果将TP和PP结合来看,每个GPU保存了本PP stage中的 个transformer块。所以每个gpu,在一个microbatch范围内需要通信的参数总量是 ,进一步化简为:

算法通信量每个(单位:参数个数)

当前LLM训练,通常是 TP/PP/DP构成的3D并行(如图11所示):

(图11. 3D并行示意图)

综上,我们可简单计算出3种并行中,计算一个batch上TP、PP、DP的通信情况。并且由于一个迭代中,有m个microbatch,所以每个GPU在一个迭代内通信情况如下:

(表4. 3D并行中per GPU通信规则)

PP的流量在Megatron-lm 第二篇paper[4]和第三篇paper[2]中都是有优化的,这里贴出了实际通信量 (单位:参数量)

在3种并行技术中:

3.3.2.2 集合通信的冗余率

1)在介绍集合通信的冗余率概念前,我们先简单看看集合通信是什么:

在3D并行中,AI框架层通过集合通信库提供的集合通信APIs,来实现多个GPU间的数据同步。如下图所示集合通信库对上提供规范的通信原语,对下是各种机间、机内的点对点通信的组合拼接。还需要考虑到底层的硬件组网拓扑、服务器内算力卡互联拓扑等。不同集群规模、硬件和互联架构下,其最优的算法和通信规则都是不同的。本小节只做相关的必要展开,后续还会有通信库相关文章做探讨。

(图12. AI/HPC应用中的软件架构)

2)集合通信为什么会有通信冗余?

在Nvidia算力卡集群中,默认使用的集合通信库是NV开源的集合通信库NCCL。NCCL中有主要有两种图算法:Flat Ring和Double Binary Tree。前者是小规模集群中使用的,后者是为了克服Ring在百卡以上范围内时延(latency)过大的问题而使用的。

我们以ring allreduce为例,简单分析下ring算法中带来的额外通信:

若有n个成员,每个成员有x的数据量要做allreduce:先将每个成员的数据x切分为n份,并通过2(n-1)个steps完成数据同步。其中每个step中每个成员都向其后一个邻居成员发送x/n的数据量。因此,每个成员共需要发送2(n-1)个x/n的数据:通过n-1个steps完成reduce-scatter,再通过n-1个steps完成allgather;

由此可以发现,在ring allreduce中每个成员即使待规约数据量仅为x,也需要每个成员总计发送x*2(n-1)/n的数据量。这里的2(n-1)/n就是集合通信算法产生的冗余率。在节点数量足够多时,使用ring算法每完成1字节的allreduce,就需要链路真实传输2字节数据。

可以简单把集合通信的数据冗余理解成“为实现有效数据同步,而不得不进行的额外数据收发”。例如图13中,nccl-test中有”算法带宽“和“总线带宽”两种指标,也是反应了集合通信冗余率的概念:算法带宽对应有效集合通信速率,总线带宽对应真实集合通信速率,总线带宽=算法带宽*集合通信冗余率。这里由于是8个GPU,所以冗余率是

(图13. NCCL-Test中反映出的通信冗余率)

3)集合通信冗余率统计:

1. 同一种软件算法中,不同的原语中的冗余率也是不同的,例如ring中allreduce是2(n-1)/n,而allgather是(n-1)/n。这里给出NCCL中Ring/Tree算法下各种集合通信的冗余率:

通信原语

冗余率

allreduce

2(n-1)/n

all2all

(n-1)/n

allgather

(n-1)/n

reducescatter

(n-1)/n

gather

(n-1)/n

sendrecv

1

broadcast

1

(表5. NCCL中Ring/Tree算法下各种集合通信的冗余率)

不同硬件架构图算法下,集合通信冗余率是有差异的:例如国产GPU卡互联方案中,主流采用full-mesh方式(缺失类NVSwitch技术),其算法对应的也是full-mesh。具体通信规则如下图所示,其allreduce中的冗余率也是2(n-1)/n:通过reduce-scatter、allgather两个步骤完成allreduce,每个步骤中的单节点都需要完成"(n-1)/n*数据量"的通信;而如参数服务器、在网计算(In-network computing)等方式,冗余率等数据还会发生变化;

(图14. 单机内4卡full-mesh时allreduce示意图)

3.4 模型显存占用估算

3.4.1 为什么要评估显存占用?

大模型的显存占用是可以预估的,这是因为大模型本身显存占用就是固定的几个部分。

当前GPU显存容量是大模型训练、推理的瓶颈:目前主流算力卡如A100,其显存只有80GB。我们在训练时,经常会遇到OOM(Out Of Memory)的报错提示。所以对模型的显存占用有一个认知和预估是非常重要的。

3.4.2 怎么评估显存占用?

模型的训练中显存占用特别大,而推理中的显存占用相对要小些。因此,我们着重分析当前LLMs的训练中的显存占用情况。

3.4.2.1 训练中显存占用分析

训练中,普遍使用混合精度训练和ADAM优化器相结合。而在这样的背景下,显存占用可分为两块:1.模型状态(model states)、2.激活值为主的其它残余状态(residual states)。

1) 模型状态

采用混合精度训练可以节省内存和提高训练速度,所以被广泛使用。而混合精度训练有 FP16-FP32 和BF16-FP32两种,两种方式中的显存占用是一致的。在分析中间激活的显存占用时,中间激活值是以16位精度格式(FP16或BF16)来保存的,每个元素占了2个字节。唯一例外的是,dropout操作的mask矩阵,每个元素只占1个字节。

标准的模型训练中,每个iteration内的流程如下。总的看下来,一个模型副本内,需要18~20Ψ的空间用于模型状态存储:

(图 15. 模型训练中每个iteration中的数据更新示意图,source from [11])

在推理过程中,由于只有前向传播,而没有反向传播,所以只需要2Ψ的空间用于存储模型状态中的模型参数部分;并且在推理中有KV-cache等技术可以节约硬件资源。另外由于推理并不依赖高精度,所以可以进一步将精度进行压缩至INT8、INT4等;

2) 残余状态

残余状态部分主要由激活值和一些临时缓存和无法再使用的碎片空间内存构成。其中,临时缓存和碎片的占用空间不大。我们着重观察激活值部分。之前有提到,激活值是在前向传播时产生的各层的输出结果,而在反向传播的链式求导过程中会被使用。由于激活值空间占用比较大,业界还有各种激活值优化技术。

首先,我们分析下全量的激活值空间占用是多少:

整个LLM模型包含了l个transformer块,每个transformer块中的激活具体可分为三部分:1.自注意力结构中的的激活值、 2.MLP中激活、3.Layer Norm中激活

1. 自注意力结构中:

因此,共计 字节;

2. MLP中:

因此,共计19bsh字节;

3. Layer Norm中:由于自注意力和MLP各自关联一个LayerNorm,每个LayerNorm需要存储输入的激活值为2bsh字节,所以共计4bsh字节;

综上,一个transformer块共需要 字节的激活值占用。整个LLM单个模型副本约有l个tranformer块,共计 字节激活值的空间占用。

接着,我们分析下基于3D并行时,激活值占用情况:

在3D并行中,每个PP stage的GPU上有l/p个transformer块。而对于每个tranformer块来说,被切分到TP上的t个GPU上来保存。在这种情况下,每个transformer块的激活值占用是和刚才分析的简单情况下的全量激活值空间占用不同。

1. 自注意力结构中:

因此,共计 字节;

2. MLP中:

因此,共计 字节;

3. Layer Norm中:发生变化。自注意力和MLP各自关联一个LayerNorm,而由于注意力结构和MLP都进行了切分成了t份,相应地LayerNorm也有t份。因此,激活值从4bsh字节变为4bsht字节;

综上,在3D并行中,一个transformer块共需要 字节的激活值占用。

我们假设目前采用的PP是pipedream的1F1B模式[13],为了保持比较低的气泡,则第一个PP stage上的micro batch个数为p(即m=p),此时第一个PP stage上每个GPU的激活值占用为:

并行中的单激活值占用

如果将l个transformer块,存在l个PP stage上。那么此时,每个PP上只有一个transformer块。每个GPU只需要存入tranformer块的输入值即可,后续所有的激活值都可以通过全激活值重计算(full activatiton recomputation)来重算。所以每个GPU只需保存2bsh 字节的数据;

3.4.2.2 显存优化技术

除了上述的显存占用外,业界还有些显存优化相关的技术。正如之前所言,这些技术通常在做显存、算力和通信之间trade off的游戏。

1)zeRO系列技术- 显存和算力的优化

zeRO技术可以看作是改进版数据并行技术:微软在Deepspeed框架中提出的zeRO系列技术,包括zeRO、zeRO Offload 及zeRO Infinity,将数据并行的显存占用、优化器参数更新的计算逐步进行卸载到CPU的内存、CPU core、NVMe磁盘上,以优化集群中数据并行中的算力和显存的资源使用。但是zeRO系列技术采取主要思想还是硬件资源间trade off:需要引入额外集群内通信来换取更小的显存占用和计算量;

zeRO是一个系列技术,主要技术概括如下:

zeRO技术中,最广泛使用的是zeRO-DP:

(图16. zeRO DP的三阶切分示意图,source from[12])

如果应用 zeRO DP中的Stage 1~3,DP中的每个模型副本中的显存占用将被进一步优化:

(表6. zeRO-DP中的模型状态的显存优化)

zeRO-DP,涉及的流量如下:

(表7. zeRO-DP中三阶段涉及的流量)

2)选择性重计算(selective activation recomputationp[2]), 会选择性地重计算部分激活值:只对生成激活值占用空间大、算力要求少的结构进行激活值重计算。相比全计算方案,整体的激活值占用、算力需求进一步优化。限于篇幅原因,不展开介绍;

3)另外,还有序列并行(sequence parallel[2])的技术,也可以优化激活值占用:通过修改张量并行的使用的通信原语,减少3D并行中重复激活值的缓存。限于篇幅原因,不展开介绍;

3.4.2.3 显存占用评估总结

综上,我们可以总结出显存占用情况:

(表8. zeRO-DP优化后的per GPU上的模型状态量)

(表9. LLM训练中的激活值占用情况)

3.5 集群扩展性估算

3.5.1 阿姆达尔定律

阿姆达尔定律(Amdahl's law)1967年由美国人阿姆达尔提出,反应了并行运算的效率提升能力。阿姆达尔定律的一种写法是:

该定律的一个关键点在于它明确了:并行系统中的加速比上限取决于系统中无法并行的串行部分,而非并行部分;例如,有个需要单线程执行20小时的任务需要使用多进程的方式来加速。如果该任务中有一个1小时的部分无法被并行执行。这个时候无论怎么去优化整个任务,最终的任务执行时间必然大于1小时;如图17中,有4个任务,每个任务的加速比的阈值上限已经确定,无论增加如何多的进程,只会无限逼近而无法超过该阈值。

(图17. 阿姆达尔定律示意图,source from [9])

那么阿姆达尔定律在当前的大模型算力集群中是怎么应用的呢?

假设在模型训练时,其网络通信占比0.05%,这个是无法通过各种并行技术进行加速,可以认为p;而GPU计算占比应该为1 - 0.05%=99.95%, 对应公式中的(1-p);其中有512个模型副本进行数据并行。那么这时加速比 = 1 / (0.05% 99.95%/512)=407.8:

由此可见,机内、机间互联传输等无法通过并行加速的任务时间,是制约整个集群能力的关键。

3.5.2 集群性能评估指标

实际上,在LLMs相关算力集群中,我们通常使用两种指标来评估集群性能:Tokens per second per GPU (单卡每秒的Tokens处理量)、FLOPs per second per GPU(单卡每秒的浮点运算量):

3.5.3 模型训练时间估算

根据阿姆达尔定律,由于整个系统中有网络通信等无法并行加速项的存在,集群中存在加速比上限。通常,训练集群中的GPU利用率通常约30%~55%之间[2,3,6]。集群训练耗时计算公式如下:

我们代入公式,进行一个试算:

已知集群中有3072 张 A100 卡,假设单卡利用率为50%;如果该集群都用于进行 1Trillion模型的预训练,整个预训练过程中喂入的 Tokens 总数为450B;集群默认开启激活值重计算;

此时套用公式后,得到预训练总时长:

4. 集群硬件资源试算和实测分析

结束了工程化参数的分析和公式推导工作,我们来看下如何通过上述信息来指导硬件选型。

4.1 硬件资源估算

在模型训练中,对算力卡的硬件性能要求较高。目前算力、通信和显存三个方向都存在瓶颈,制约了集群整体表现。并且,不同的软件配置下,整个集群中的显存、通信和算力的压力都是不同的。

4.1.1 集群规模估算

目前,绝大部分卡在16位精度(FP16或BF16)下算力水平都在300TFLOPS左右。所以我们以此为算力卡基准算力,分析当前常见的几种规模的模型需要的卡数:

1)GPT3-175B

我们代入paper[4]中175B模型训练任务的预设条件:”使用300B tokens进行训练,总耗时23天“。而此时,我们以单卡300TLOPS作为业界算力卡基准,分析预设条件中需要的卡数:

算力卡数算力总需求单卡峰值性能单卡利用率时间总数单参数量的算力需求模型参数量单卡峰值性能单卡利用率时间算力卡数未开启激活值重计算张

2)LLaMa-65B

在论文[15]给出了llama 65B模型的训练情况:”使用了1.4T的tokens,训练用时约21天”。我们代入单卡300TFLOPS的算力,分析预设条件中需要的卡数:

算力卡数未开启激活值重计算张

3) LLaMa-13B

[15]中给出了13B的模型是1T的tokens来训练,在此基础上如果我们假设需要15天来完成训练,则需要的算力卡数量为:

算力卡数未开启激活值重计算张

4.1.2 显存占用估算

根据3.4.2.3 小结中的计算公式,我们可以计算出显存的占用情况。整个GPU显存占用可以份为2个部分:

模型状态缓存无优化模型参数量模型状态缓存模型参数量模型状态缓存模型参数量模型状态缓存模型参数量

激活值无重计算激活值选择性重计算序列并行激活值全激活值重计算

整个GPU的显存占用,等于模型状态和激活值缓存。由于当前zeRO3对于系统的通信损耗过大,目前训练主流还是zeRO1或zeRO2。我们也将试算单卡显存占用2种典型值:1.有zeRO1的3D并行 选择性重新计算 序列并行;2.无zeRO的3D并行 无激活值重计算:

我们对于各种模型进行配置敲定:

1)llama 13B

在上一小节中,已经计算出llama 13B在使用1T tokens、用15天完成训练时,需要约300TFLOPS的算力卡约422张。根据[16],13B模型的常见并行度是t=2,p=1。为了3D并行度和机器数量凑整,我们令GPU总数为424: t=2、p=1、d=212。此外,B需要是b*d倍数,所以我们可以令B=1272、b=1、s=4096;

2)llama 65B:

上一小节中 ,已经计算出65B的模型在使用1.4T tokens、用21天完成训练时,使用2113张卡。而目前65B一个推荐并行度是t=4,p=4。为了3D并行度和机器数量凑整,我们令GPU总数为2112:t=4、p=4、d=132。再根据限制条件B需要是b*d倍数,设B=2112、b=1、s=2048;

3)GPT-175B

根据[2,4],目前175B中的TP并行度、PP并行度通常设为(t=8,p=8)或者(t=8,p=16)。我们以(t=8,p=8)为例,原定总卡数为1110,为3D并行和机器数量凑整,可设整个集群卡数为1152张(t=8,p=8,d=18)。B=1152、s=2048:由于流线型中micro-batch份数需要为整数,所以我们令B为1152;

此外,我们给出机器数量不变时,3种并行度变化下的显存计算结果进行参考 : 由于3种并行中dp通信开销最少,优先考虑高dp并行度下进行tp和pp的变换,结果见表10:

(表10. LLMs显存占用试算表)

由上表可见,

4.1.3 通信需求预估

4.1.3.1 通信瓶颈估算方法

一、确定LLM中的通信瓶颈

在LLM训练集群中,每个GPU持续在做3件事情:1.从关联的GPU接收到运算所需要的数据、2.使用数据,进行计算 、3.计算生成的数据发送给下一个关联GPU。从这个角度来看,每个GPU都像是工厂流水线上的自动化机器人:每个机器人需要从相邻机器人处拿到待处理工件,加工后并把完成工件传递给下一个机器人继续加工。所以,每个机器人获取工件的速度需跟得上其加工速度,否则就会出现异常等待、最终导致工厂整体效率下降;

回到GPU本身,此时可以发现存在类似道理:如果每个GPU通信和计算之间可以完全同步,那么集群是最高效的。 但是,真的可以做到通信和计算两种资源之间的同步进行吗?

答案是不能的。观察下TP、PP和DP三种技术,分析3D并行中的真实的计算和通信流程:

1)DP

DP中的梯度同步是发生在每次迭代的前后向传播后进行的,如图18。也就是说标准情况下,DP中的通信需要前后向计算都结束后才能开始。

(图18. 简单数据并行示意图, source from [13])

但实际上,根据神经网络的链式求导原理:在反向传播开始时,最后一层往前的梯度数据就在同时进行了。所以在DP上有overlap的优化(如图19):梯度allreduce几乎和反向传播同时进行。那么只要保证梯度allreduce的时间能够被反向传播时间所包含,整个集群就是高效无阻塞的。目前非大模型领域,已有些类似的优化方式(如[14]);但大模型领域若开启3D并行后情况比较复杂:涉及到多个物理副本和逻辑副本,所以目前没有采用这种overlap方式。当前,DP维度的梯度同步,主要还是等前/后向传输完毕后才进行的。

因此,我们在3D并行中,通常对于DP上梯度allreduce的通信可以要求固定占比对应计算时间(例如5%)。如果这个比率过大,根据阿姆达尔定律,整个集群的利用率将变低。

(图19. 简单数据并行中的通信overlap示意图, source from [14])

2)PP

PP中所产生的通信理论上是可以overlap的。如图8所示:

而表4展示了,每个GPU在PP 上引入的通信量为一次迭代前、后向分别传输 参数 。只要带宽足够,这部分通信时间能够被前、后向传播时间cover,整个集群就是高效无拥塞的;但并不是所有的PP都是采用的overlap的方式来进行通信传输的,我们还需要保证每次通信时间不超过相关计算时间的一个时间占比 ,以保证非overlap技术方案中的高效性;

3)TP

TP中通信和计算无法overlap。attention和mlp中都是需要等待计算结束后,才能通过通信原语进行数据同步。所以我们可以设定每次通信时间不超过相关计算时间的期望比值,如5% ;

综上,我们可以发现:

二、计算各种并行中的带宽需求

结合表3, 我们可以估算训练时,perGPU上每个迭代中的通信情况:

(表11. LLM训练中per GPU通信量)

注:

应用sequence parallel(SP)时,在NV集群中真实通信量保持不变:

根据表5得,在NCCL中Ring/Tree算法中关于t个节点的allreduce通信冗余率为 ,关于t的allgather和reduce-scatter各为 。

因此从通信量上看,在NV集群中开SP前后的TP上通信量不变。

另外,我们还可以进一步分析各种并行度下的带宽需求和参数间的关系:

  1. TP:

带宽需求数据量通信耗时通信量单单计算耗时经验的期望比率单次通信量次数数量单模型利用率单卡算力性能期望耗时占比至利用率单卡算力性能期望耗时占比至利用率单卡算力性能期望耗时占比其中为常数

可以发现,TP中通信带宽压力的公式和t成正比,和h成反比:

1)在模型变大时,h是逐渐变大的: 当张量并行度t固定时,TP带宽要求是随模型规模变大逐步变小的。例如,175B和1T模型训练中,一般都使用t=8,而1T模型中的h更大,所以1T模型的TP上的带宽压力比175B的要小;

2)在现实中,我们在模型中通常使用尽量小的t和p。这是因为TP和PP中的通信开销较为昂贵:TP通信频率最大,TP的通信频率是和b、s、h、l都有关系的。PP通信和b、s、h相关。而DP通信只和s、h有关;另外,由于TP间集合通信在通信量不变的情况下,NCCL集合通信冗余率是 ,在t越小时真实的通信量越小,所以t会倾向取小。例如在175B、65B、13B模型上,TP并行度(t)通常分别为8、4和2。

2. PP:

带宽需求数据量通信耗时通信量单单计算耗时经验的期望比率单次通信量次数数量单模型利用率单卡算力性能期望耗时占比至利用率单卡算力性能期望耗时占比其中为常数

可以发现,PP带宽压力和每个PP stage上存在的transformer block数(l/p)成反比,和h也成反比;当模型越大时,若每个PP stage上的blocks数量不变,那么其通信压力不变;

3. DP:

带宽需求数据量通信耗时通信量单迭代单迭代计算耗时经验的期望比率单次通信量次数数量单模型利用率单卡算力性能期望耗时占比至利用率单卡算力性能期望耗时占比其中为常数

可以发现,DP并行带宽压力和数据并行度d成正比,和一个迭代喂入整个系统的tokens数量(Bs)成反比;

上述公式中,TP和PP中带宽要求和模型规模成反比的原因:通信和计算两个任务中,虽然在模型尺寸变大时两者都会发生变化。但是两者变化的速度是不同的:模型规模越大时,算力需求提升的速度(和h的平方成正比)比通信传输量(和h成正比)提升的速度更快。

4.1.3.2 通信带宽估算

根据表4的公式,进行1次迭代时的各种通信相关的数据计算方法如下:

每个microbatch的计算耗时为:

数量单个单个模型副本算力需求单个模型副本数量单卡算力单卡利用率至单个模型参数量单卡算力单卡利用率

一次迭代中的计算耗时为:

数量

TP的通信量和通信频率的计算:

单次通信量通信次数单通信次数单迭代通信次数单

假设TP通信和计算的时间消耗比重是5%。由于已经获得出TP中的计算耗时和TP通信量,则可预估出所需单GPU对外通信带宽。推导过程如下(下面公式是单个迭代的计算方法,使用单个microbatch的方式是类似的):

带宽需求数据量耗时通信量单迭代计算期望计算和通信占比通信量单迭代计算

单次通信量通信次数单个通信次数单个迭代带宽需求通信总量期望通信时间单次通信量通信次数单个迭代计算期望计算和通信占比单次通信量单个迭代中的通信次数计算

单次通信量通信次数单个迭代带宽需求通信总量期望通信时间单次通信量通信次数单个迭代计算耗时期望计算和通信占比

使用上述公式,对于4.1.2预设的各种模型和其配置进行通信带宽试算:

(表12. LLM训练中通信带宽试算表)

在机器硬件参数确定时,我们还可以通过上述公式去反推出合适的并行度参数。

5. 小结

本篇文章主要展示了LLM中硬件资源的估算方法和硬件瓶颈的理论推导,让大家了解大模型硬件选型等工作是如何开展的。我们发现了大模型中算力、带宽和显存三者都很重要,缺一不可: 例如只增加单卡算力时,由于通信瓶颈存在并不会带来集群性能的相应提升;此外,硬件只决定了集群性能的理论上限、集群中软件和组网等还会导致上限的进一步收缩,因此集群组网设计和拥塞控制算法配置、集群通信库算法等都是非常重要的。有机会的话,我们将于后续的系列文章继续介绍。

参考

查看全文
大家还看了
也许喜欢
更多游戏

Copyright © 2024 妖气游戏网 www.17u1u.com All Rights Reserved