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

1)  算力瓶颈：

• 进行大模型训练时，每输入一个token，整个模型中的每个参数上要进行6~8次的浮点数运算；而若以300万tokens的数据集来训练一个175B的GPT3模型，那么使用单张理论峰值算力为312TFLOPS(FP16)的A100 GPU，即使不考虑单卡在集群中的性能折损，也至少需要32年才能完成。这是无法被接受的；

2）显存瓶颈：

• 大模型训练中，LLM训练时的显存压力非常大。微软&英伟达[1]给出了模型状态显存占用公式：单个模型副本中每个参数量大约需要20倍于自身大小的空间占用(基于混合精度[10] & ADAM优化器的SOTA技术来完成训练)。以175B模型训练为例，因为单个模型参数量达到175Billion，则至少需要3.5TB的显存空间占用。而目前主流算力卡如A100、H100只有80GB的显存空间，所以至少需要44块80GB显存的GPU才能塞下一个模型副本。并且由于算力需求巨大，整个系统会有多个模型副本，总的显存占用是很夸张的。此外，模型训练时，还需要保存激活值用于反向传播。激活值还将额外占用很大一部分空间。虽然有优化激活值空间占用的技术，如激活值重计算(Activation Recomputation)等，但是会引入额外的通信量和计算量；
• 模型推理中的显存压力相对小些，其中一个原因是推理中计算对精度要求不像训练那般高(训练中的优化器状态更新需要在至少32位精度下进行)。另外，由于推理和训练表现是关联的，只需保证训练出的模型的优质，即使推理时缩减参数精度也不会有很大影响。但即便如此，由于LLM通常拥有足够大的参数量，推理只需1~2倍于模型参数的空间占用也是惊人的；

3）通信瓶颈：

• 由于大模型训练&推理通常都不能在单卡进行，所以我们需要多算力卡构成集群，不可避免的就存在服务器机内和机间通信。由于大模型本身的结构宽度、激活值尺寸都非常大，根据张量并行、流水线并行等通信规则，机间和机内通信压力都会很大。另外，虽然当前使用多种并行技术将集群利用率保持的比较好，但这是榨干的各种硬件互联带宽资源换来的；
• 整个集群的通信，不在是简单的点对点(send&recv)，而是涉及到多个节点间的集合通信(collective communication)概念。集合通信虽然给AI框架封装提供了高效易用的APIs接口、并基于硬件互联架构优化了底层数据流量，但会引入额外的通信冗余：例如使用Nvidia的通信库NCCL进行Allreduce时，每处理1Byte的数据量就需要约2Byte的传输。而我们整个集群在训练和推理阶段，充斥着高频且大量的集合通信操作，所以会带来非常严重且不可避免的数据传输损耗；
• 由于算力和显存都存在瓶颈，业界诞生和迭代了很多硬件资源优化技术，如zeRO、选择性激活值重计算等。但其中绝大多数，只是在做不同硬件资源间的”trade off”：用通信换算力、用通信换显存等。这也导致整个集群中通信压力进一步增大；

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

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

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

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

1）首先，多头自注意力结构共包含   参数量：1）该结构一共有 a 个头。每一个头都拥有一个self-attention结构，其中包含   三个矩阵，用于生成 Q、K、V 三个矩阵。而   每个矩阵形状是(h，h/a)。因此：  三种矩阵的参数量是  。此外每种矩阵还包含了h个偏置。总计；2）另外，后面有个 WO矩阵形状是(h，h)，参数量是  ，还包含了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结构也基本无参数量；

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

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

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

首先，我们可以发现，矩阵乘法(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)。

上述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位精度进行，整个任务的计算效率非常高。

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

• 模型的训练，就是通过多次迭代(iteration)计算完成的。每次迭代包含前向传播、反向传播两个步骤。1）前向传播中，输入值和当前层的模型参数做矩阵乘法，其产生的结果称为激活值。激活值再作为后一层的输入值和后一层参数做矩阵乘法，所以整个前向传播就可认为主要是一系列链式的矩阵乘法；2） 而反向传播，是从最后一层开始，应用链式法则反向逐步进行激活值偏导数和梯度的计算。其算力需求大约是前向传播的两倍[4]；3）前向传播过程中产生的激活值是需要被暂存的，并且在反向传播中被使用。而大模型场景下，产生的的激活值总量通常非常大。业界通常会采用激活值重计算进行显存优化：在反向传播开始前重新进行前向传播计算，如果采用的是全激活值重计算(full activation recomputation)，则额外引入了和前向传播同量的计算工作量；
• 模型的推理，相当于就是一个前向传播的过程，经过一系列链式的计算，最终生成预测结果的tokens。

我们令模型需要浮点计算次数(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）注意力结构中，

• 每个head中都发生了前一层传入的激活值([,s,h]的矩阵)和  中的每个([h,h/a]矩阵)相乘的运算，这部分共涉及算力为   FLOPs。 而由于有a个head，所以整个结构需要的   FLOPs；
• 另外，有些非标准矩阵乘的计算，共有  FLOPs：在每个head中，得到KQV三个矩阵结果后，需要进行  ，以实现同个sequence内tokens间的Q与K点积运算、缩放。Q的形状[  ,s,h/a]，  的形状为[  ,h/a,s]，涉及个[s,h/a]*[h/a,s]矩阵乘。共有  的算力需求。由于共有a个头，所以共   FLOPs；之后，将softmax结果与V进行element-wise相乘，前者维度为[  ,s,s]，后者维度为  ,s,s/a]，需要   FLOPs。由于共有a个头，所以共  FLOPs；
• 此外，WO是一个[h,h]的矩阵。之前的激活值可以认为是一个[,s,h]的矩阵。这里矩阵乘法为[b,s,h]*[h,h]，共有   FLOPs；

单个attention共包含   FLOPs。

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

• 激活值和第一个线性层相乘：[B,s,h]*[h,4h] 共   FLOPs；
• 激活值和第二个线性层相乘：[B,s,4h]*[4h,h] 共   FLOPs；
• 激活层为约0FLOPs：是确定的数学缩放公式，可认为无需矩阵乘；

单个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；

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

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

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

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

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

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

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

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

^{(表3. 参数含义表）}

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

一、  数据并行

1）简单数据并行

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

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

• 同步模式，如NCCL中的Ring、Tree算法等，是需要等待所有节点上的完成数据同步后，才会开始下一次数据并行；
• 异步模式，例如参数服务器，是无需等待所有节点都完成数据同步。每个节点只需要保证获得最新的参数副本即可；

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

• 优点：1.具备通用性，不关心模型的具体结构，实现比较简单；2.在多种并行技术中，DP拥有最少的数据通信频率和通信量：每次迭代仅做一次梯度同步，是在无硬件瓶颈下的首要并行选择；
• 缺点：在数据并行中的每个模型的副本中，都要缓存一个完整的模型状态数据：高达约20Ψ个数据；

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

二、模型并行

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

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

• 流水线并行(PP)，指的是把LLM中的l个transformer块看成一个序列。通常把序列进行几等份的均匀切分，每个算力节点保存一个由几个transformer块串联成的一段结构。对应术语做层间并行(Inter-layer parallelism)
• 张量并行(TP), 指的是把每个transformer块进行更细化的拆解：由于LLM模型主要是多个重复的transformer块构成。把每个tranformer块的结构进行拆解，这就是TP的主要思路，对应的术语叫做层内并行(Intra-layer parallelism)；

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

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])编排的更加复杂些，但也能获得更小的气泡。但两者其基本思想类似,所以就不展开讨论了。

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

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

• 前向传播时前一个PP stage将本GPU上最后一个输出的激活值(bsh个参数量)传播给下一个PP stage对应的GPU，用于后续链式计算；
• 反向传播时后一个PP stage将激活值偏导数(bsh个参数量)传递给上个PP stage，也用于链式计算；

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

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

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

• 多头的attention结构可以按head进行拆分，每个head都可以放到不同的GPU上；接着，后续衔接的WO矩阵可以按行拆分到各个GPU上，以实现和当前头内KQV投影的矩阵乘；最终每个张量并行的GPU上都会获得一份[b,s,h]的激活值，需要对于该激活值进行allreduce同步(对应图9中的绿色矩形”g”)。由于训练中前向传播和反向传播都是链式的，而反向传播也需要对于激活值的偏导数进行allreduce(对应图9中的绿色矩形”f”)。所以每个microbatch，涉及2次allreduce；

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

• d*p*t个gpu构成了整个系统: 系统中，每p*t个GPU构成了一个物理模型副本，共有d个这样的物理模型副本。每个物理副本中，还有因pipeline机制引入的多个逻辑模型副本；
• 物理模型副本间采用allreduce进行梯度同步，同步频率是每次迭代同步一次，总量上是关于模型梯度(共Ψ个参数量)。通常, 不同物理副本上的同号GPU使用各自的RDMA网卡进行同步，每个GPU的通信量为  个参数；
• 而物理副本内，有  个逻辑副本，逻辑副本间也是关于梯度的allreduce同步。但由于逻辑副本模型分片都是和GPU绑定的，所以每个GPU上逻辑副本间的梯度同步无实质通信量；
• 而逻辑副本内，GPU之间还有TP和PP的通信量，这里就不赘述了;

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

• TP中的通信频率最大、通信量最高，因为它和输入的tokens数量、模型宽度、transformer层数都有关。所以TP一般限定在机内，并由GPU间高速互联(如NVLink)进行；
• PP的通信仍然和输入tokens数量、模型宽度都有关，但不和transformer层数有关，所以PP中的通信量和频率比TP小；
• DP具有最少的通信量和通信频率，它和输入的tokens数量无关；

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

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

在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，所以冗余率是

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

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

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

1)  模型状态

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

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

• 先将前一层传入的激活值作为输入值，结合模型的保存的16位精度的参数(2Ψ)来进行前向、反向传播后，得到梯度数据(2Ψ)：在这个过程中，分别需要2Ψ显存存储16位精度的参数、梯度数据；
• 之后，使用计算出梯度数据，进行优化器状态状态更新，计算得到包括计算一阶动量(momentum)、二阶动量(varience)：一阶动量和二阶动量都是32位精度缓存，分别需要4Ψ的空间；
• 接着，使用一阶、二阶动量来更新参数。在这个过程中，需要先将需要一个4Ψ的梯度 buffer 来存储参数即将变化的差值信息(如果此时将之前16位精度的梯度缓存给释放，此处只需要额外增加2Ψ的空间)；此外，使用之前的参数减去变化差值得到更新后的参数，这还需要4Ψ的空间用于在32位精度更新参数。
• 最后，得到的更新后的参数，由fp32转fp16。并且将fp16的参数用于下一个iteration的计算；

2)  残余状态

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

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

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

1.  自注意力结构中：

• QKV矩阵乘：共享bsh的输入激活值，对应2bsh 字节；
•   矩阵乘：由于存储了Q和K两个矩阵，而Q和K的形状都是(b,s,h)。所以需要存储2bsh的参数，共计4bsh字节；
• Softmax：需要存储  矩阵乘的结果作为softmax的输入，而  的形状为(b,a,s,s)。所以需要  字节；
• Softmax dropout：仅需要保存一个mask矩阵，其结构  和一致，所以需要  字节；
• Score*V: 分别需要存储score和V，score对应  字节，V对应2bsh字节；
• 结果和WO矩阵乘：需要保存其输入，大小为 2bsh字节；
• dropout：需要保存mask矩阵，大小为 bsh字节;

因此，共计  字节；

2.  MLP中：

• 第一个线性层：保存其输入的激活值，需要2bsh字节；
• 激活函数：保存其输入的激活值，需要8bsh字节；
• 第二个线性层：保存其输入的激活值，需要8bsh字节；
• dropout：需要保存mask矩阵，大小为 bsh字节;

因此，共计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.  自注意力结构中：

• QKV矩阵乘：发生变化。原本多头间共享bsh个输入值(2bsh字节数据) 。 现在做TP产生了t个切分，每个切分上都得保存一份输入值。所以激活值从2bsh字节变为2bsht字节；
• 其余部分保持不变；
• WO后的dropout：发生变化。有t个切分，每个切分都得保存mask矩阵，总大小由bsh字节变为bsht字节;

因此，共计  字节；

2.  MLP中：

• 第一个线性层：发生变化。该矩阵被认为按列切分成t个子矩阵。所以每个子矩阵都得保存一份输入的激活值，激活值从bsht字节变成了2bsht字节；
• 其余部分保持不变
• dropout：发生变化。需要保存t个mask矩阵，大小为由bsh字节变为bsht字节;

因此，共计   字节；

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的激活值占用为：

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

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

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

zeRO是一个系列技术，主要技术概括如下：

• zeRO 提出将zeRO-DP和 zeRO-R技术。zeRO-dp负责将model states进行分割并且存储到各个进行数据并行的gpu上，包含zeRO stage1~3；zeRO-R进一步优化residual states的显存占用;(部分存储卸载，通信换内存)
• zeRO-offload，提出将模型的model-states计算和更新卸载到cpu中进行(部分计算和存储卸载到CPU和CPU memory中，通信换内存)；
• zeRO-infinity，进一步盯上了服务器上的NVMe磁盘资源。将显存占用卸载到NVMe上进行；

zeRO技术中，最广泛使用的是zeRO-DP：

• zeRO-DP 是将显存占用所需要的model states进行切分，然后每个数据并行的副本仅仅保存其中的一部分。如图6所示， 具体分为三种方案：1.stage1 是仅仅切分优化器状态()；2.stage2是切分优化器状态和梯度(  )；3.stage3 是切分优化器状态、梯度、参数(  )；
• 通常大部分的 AI 框架都支持和默认开启了 类zeRO 的数据并行优化机制，例如 Pytorch 中的 FSDP(Fully Sharded Data Parallel)， 就是 zeRO3的思想；

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

综上，我们可以总结出显存占用情况：

• 模型状态(model states ，约18～20Ψ)：1.参数(2Ψ)、2.梯度(2Ψ)、3.优化器状态(一阶动量 、二阶动量、梯度buffer、参数buffer， 共计16Ψ)；如果应用zeRO-DP技术后，这部分显存占用会发生如表8的变化:

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

• 残差状态(residual states)中主要是激活值占用。激活值占用情况如下表9：

• 其中   表示的是一个百分比，代表的是整个任务的理论加速比 ；
• s表示可以通过并技术加速的倍率，例如在训练中对应GPU总数;
• p表示可以通过并行技术进行加速的任务时间比重；
• (1-p)表示无法通过并行技术进行加速的任务时间比重；

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

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

• 这表示相比单卡，如果可并行的部分进行提速512倍的话，整个系统的速度最多能加速到原来的407.8倍；
• 此时，如果比较的512张卡受网络制约的情况和 512张卡无制约的情况，最终得到的加速比的407.8/512 =79.6% 。即所有GPU在集群中至多发挥79.6%的单卡的算力性能；

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

• Tokens per second per GPU ：通过统计一段时间内，模型吃入或生成的Tokens 数量。并且除以集群中的GPU卡数和时间，得到单卡的tokens处理速度；
• FLOPs  per second per GPU：由于整个模型的算力总量可以预估，此时只需要再获得实际进行模型训练/推理的时间，就可以得到实际的单卡每秒的浮点运算量；

• 未应用激活值重新计算时：

   

• 应用了全激活值重新计算时：

   

我们代入公式，进行一个试算：

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

此时套用公式后，得到预训练总时长：

1）GPT3-175B

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

• 之前提到集群中的GPU利用率通常约30%～55%之间，我们设此时集群中的单卡利用率在50%；
• 根据3.5.3节的公式，易得单GPU上的算力性能计算公式为：

• 由于训练时激活值占用显存空间过大，通常为了节约内存还会开启激活值重计算的技术：1.全激活值重计算虽然能获得特别好的激活值优化表现，但是整体算力需求上升33%; 2.选择性激活值重计算结合序列并行能保持算力和显存之间较好的trade off：在降低5倍激活值占用[2]基础上，只需要提高约5%[2]以内的额外算力开销。因此，我们以训练开启选择性重计算进行预估，需要1110张GPU卡（1057*1.05)；

2）LLaMa-65B

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

• 假设此时集群中的单卡利用率在50%。则算力卡数即可如下推导

• 我们依然假设集群采用选择性激活值重计算和序列并行，那么需要的GPU数量为2113张(2012*1.05)；

3)  LLaMa-13B

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

• 假设集群中单卡利用率在50%，则卡数计算方法如下

• 我们依然假设集群采用选择性激活值重计算和序列并行，所以13B模型需要约422张卡(402*1.05)；

• 1. 单GPU模型状态：

• 2. 单GPU激活值缓存：

整个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：

• 要支持相应规模的模型训练，以常见的zeRO1+选择性重计算为例，单张算力卡的显存占用在30-80GB之间。表格中有些并行度下的显存占用较小，但不具备参考意义，是因为这些并行度下已经达到了通信瓶颈，我们将在4.1.3中的表格中详细说明。此外，由于不同的AI框架的显存优化效果不同，在增大batch size时实际显存占用可能要比上表要大；
• 需要额外说明的是，在模型规模上升时，单卡的显存需求并没有显著上升：因为单个模型副本的显存占用虽然增大了，但是相应地我们使用了更大的TP、PP的并行度 ，这使得每个GPU上的模型参数量和其显存占用并没有明显上升；
• 此外，若在GPU显存大小确定时，我们还可以通过本节的显存计算公式反推出合适的并行度参数。限于篇幅，不做展开；

一、确定LLM中的通信瓶颈

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

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

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

1）DP

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

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

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

• 前向传播中，每个PP stage的GPU上产生的激活值数据可以在下一个计算流程(即反向传播时)发送给下一个PP stage的GPU；
• 反向传播中，每个PP stage的GPU上产生的激活值偏导数据可以在下一个计算流程(即前向传播时)发送给上一个PP stage的GPU；

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

3）TP

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

综上，我们可以发现：

• 对于确定可以overlap到计算时间段中的通信任务，我们仅需要其耗时小于计算的耗时。从而确定最小的通信带宽要求；
• 对于无法overlap到计算中的通信，我们可以设定一个预期的计算通信耗时占比，例如5%。从而预估出最小的通信带宽要求；

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

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

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次迭代时的各种通信相关的数据计算方法如下：

• 耗时计算：

每个microbatch的计算耗时为：

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

• TP通信计算：

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

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

• PP相关信息计算：

• DP相关信息计算：

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

• TP通信一般对应如NVLink的机内高速互联：由上表可见，运行经典规模的三种大模型，机内互联带宽一般需要在单向100-300GB/s之间。而目前国内可买到的NV GPU机内互联带宽最高为单向200GB/s(如A800、H800)，主流国产算力卡的机内互联带宽单向最大也不会超过200GB/s；前面提到，我们可以通过调低TP并行度来减少机内互联带宽压力，但其他两种并行度也需要协同调整，这可能会带来其他的集群影响，需要结合表10和表12综合评估。
• DP的通信一般走机间RDMA网络通信：由上表可见，需要每个GPU独享100-400Gb/s的RDMA网卡带宽；前面提到，DP的带宽和B、s、d都有关系，我们可以尝试调大Bs或者调小数据并行度d来减少网络带宽需求；
• PP通信使用的物理链路和并行度有关，和机内通信和机间通信都有关系。但从上表可以看出，PP的通信带宽相比DP、PP都是比较小的，对硬件没有额外的压力；

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