大模型训练到底需要多少算力?Qwen2-72B训练资源全解析(建议收藏)
文章摘要: 本文探讨了大模型预训练的算力需求问题,重点分析了Qwen2-72B模型的训练时间和资源计算。研究提出了算力需求公式3T(2.6e6s + 2P),其中T为数据集token数,P为模型参数量,s表示序列长度。结果显示,7T tokens数据集在6000张A100上训练一个epoch约需30天。通过分析模型架构和计算过程,发现大模型算力主要消耗在矩阵乘法(占99%),其中反向传播是正向的2
文章分析了训练大模型(特别是Qwen2-72B)所需的算力资源,提出了计算公式3T(2.6e6s + 2P),其中T为数据集token数量,P为模型参数量,s表示序列长度。大模型计算量主要与矩阵乘法相关,反向传播是正向的2倍。序列长度对算力需求影响较小,batch size超过阈值后对训练时间无影响。通过实际案例验证了公式的准确性。
背景
让我们想象一个场景,假如说某一天,上面哪个大老板突然拉着你进一个会议,会上一群人问你:假如给你一个千卡集群,让你训一个qwen/百灵/凤凰/chatglm/chatgpt出来,能不能搞?有什么困难?需要多久?
此时肯定是结论先行:能训。此时此刻,非我莫属。
但第二个问题怎么回答呢?需要一堆数据,需要多少T,需要多少人力?需要一个工程组来搞定千卡架构的问题,需要一些算法hc,招人嘛。
那就剩下第三个问题了,需要多久才能训好一个大模型?日常算法训练一般是直接拉上去先跑着,从tqdm进度条和loss曲线大概预估下。但开着会,总不能说先让我把数据,工程组准备好,我先跑几个step试试?
本文就是回答这个核心问题。预训练一个模型,需要多久。
结论
1.预训练一个qwen2-72b,给定7T tokens数据集,6000张A100,一个完整epoch需要最多30天。训练语料的长度是在预训练最后阶段才从4096拓展到32768(但笔者没找到这个时间点)。所以本文预估的算力需求会有一定程度高估,但不会超过1.6倍。
2.计算量需求公式为3*T(2.6e6*s + 2P),其中T为数据集token数量,P为模型参数量,s表示序列长度。在序列长度较短时退化为6TP。若使用了全部重计算技术,则系数由3变成4。
3.大模型计算量只和“矩阵乘法”有关。且反向传播过程是正向的2倍。不同优化器影响不大。
4.attention对seq长度的平方复杂度,拉到32768长度对总算力需求也就是增加0.6倍。
5.batch size对计算量没有影响。在超过某个阈值后对训练时间没有影响。
正文
基础概念科普:
FLOPS
定义:floating point operations per second,每秒浮点运算次数。即常规理解中,硬件的性能(计算速度/算力)。
注1:GPU的算力正常是打不满的。涉及到各种框架、并行、木桶原理、通信、调度、内存的概念。可以出100道面试八股文,淹死一整个算法工程团队。正常记住几个结论对算法就够用了:
A100,单卡单精度,利用率MFU一般在25到75%之间(FlashAttention2能拉到上限),取个居中的50%,约等于300 T FLOPS。
H100,算力是A100的三倍多一点,利用率MFU需要最新出的FlashhAttention3才能拉到上限,一般可取1000 TFLOPS。
注2:即使是同一个gpu,对不同精度的运算,性能也是不一样的。这个涉及到硬件设计的架构实现,以及各种精度运算的硬件实现。这里不赘述。
详细原因可参考:
不同产品的计算能力:https://developer.nvidia.com/cuda-gpus
计算能力解释:https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capability-9-0
英伟达A100与H100以及利用NVLink技术将两块H100互连的GPU在不同精度下FLOPS的对比。这里一个T就是10^12。
FLOPs
定义:floating point operations,浮点运算数量,即常规理解中,训练一个大模型需要的算力。可以用来衡量算法/模型的复杂度。乘法和加法混同看待。
1 MFLOPS(megaFLOPS)等于每秒一百万(=10^6)次的浮点运算。
1 GFLOPS = 10^3 MFLOPS(gigaFLOPS)等于每秒十亿(=10^9)次的浮点运算。
1 TFLOPS = 10^3 GFLOPS(teraFLOPS)等于每秒一万亿(=10^12)次的浮点运算,(1太拉)。
1 PFLOPS = 10^3 TFLOPS(petaFLOPS)等于每秒一千万亿(=10^15)次的浮点运算。
1 EFLOPS = 10^3 PFLOPS(exaFLOPS)等于每秒一百京(=10^18)次的浮点运算。
1 ZFLOPS = 10^3 EFLOPS(zettaFLOPS)等于每秒十万京(=10^21)次的浮点运算。
也是本文的主要内容。
MACs
定义:Multiply-Accumulate Operations,乘法加法累积操作次数。是深度学习领域最常见的计算的一种抽象。即将两个数相乘,并将乘积累加到一个累加器上。
也是描述(训练一个大模型需要的)算力的一种单位。
按照定义,1MACs ≈ 2FLOPs。
MACs用的不算多。根本原因在于正常大模型计算中,乘法和加法就是一比一的。导致没必要单独算这个。
一比一的原因可以继续看下面的分析。
硬件上矩阵乘法的算力需求
假设我们有个矩阵A : a1 * a2,有个矩阵B : b1 * b2。我们需要计算C = A * B
首先根据定义,a2 = b1,不妨设他们等于h。最终输出的矩阵C : a1 * b2。
如果记不得矩阵乘法的定义,可以参考这张图:
从C反推,a1 * b2个元素中。每一个元素都需要经历h次乘法,和h次加法。即,2h FLOPs = 1hMACs。
注:这里之所以不是h-1次加法,是因为硬件计算加法的本质是需要放到累加器里。所以哪怕是初始第一次乘法结果也得做一次加法。
即,A * B的矩阵计算,需要的算力为 2 * h * a1 * b2 FLOPs,即 2 * h * 输出矩阵参数量 FLOPs。
大模型FLOPs计算
这里先摆一个qwen2-72b模型架构图。
放一下参数:
{
放一下模型运算py文件:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2/modeling_qwen2.py
前向计算过程
结论
大模型的算力需求基本只看矩阵乘法。
抽取出需要的参数:
batch size:用户自己定义,假设是4
seq length序列长度:取个比较大的长度32768
hidden_size隐藏层大小:8192
num_hidden_layers:80
vocab_size词表大小:152064
Embedding层(参数量占比1.7%,算力需求占比0%)
输入:[batch size, seq length]
输出:[batch size, seq length, hidden size]
这层需要将输入的token序列映射为对应的embedding序列。
即,需要look up每一个输入token在词表中的embedding。
这里会涉及到一些position embedding的计算,例如输入序列长度超过:
max_position_embeddings
会临时计算新的position embedding。不超过就直接取计算过的缓存等等。
但因为计算量实在太小,可忽略不计。
Transformer层(参数量占比96.6%,计算量占99%)
单个Transformer主要包括一个Attention块和一个FFN块,还有其他杂项,分别计算。
单个Attention块(参数量占比16%,计算量占48%)
输入:[batch size, seq length, hidden size]
输出:[batch size, seq length, hidden size]
挨个拆解步骤:
1.将输入分别映射到三个QKV矩阵上(注意由于GQA,kv矩阵小一截)。
a.Q:q_proj * hidden_states,计算量 2 * hidden_size * num_heads * head_dim * batch size * seq length = 17,592,186,044,416 ≈ 17.6 TFLOPs;
b.K计算量减少num_attention_heads / num_key_value_heads倍,即2.2 TFLOPs;
c.V同k,2.2 TFLOPs;
d.注:这里涉及到kv cache的问题。为了简化讨论,以空cache状态计算量为参考。可以理解为忽略了kv cache技术。若慢cache,则kv矩阵计算量几乎为0。
2.应用旋转向量。计算量很小。
3.K、V矩阵拓展到num_attention_heads个头,放大num_attention_heads / num_key_value_heads倍。计算量很小。
4.Q * K:
a.此时Q矩阵和K矩阵形状一样,大小都是[batch size, seq length, hidden size];
b.上面这条算的是错误的。为什么错,多头注意力的多头计算就在这里要展开了;
c.此时Q矩阵和K矩阵形状一样,大小都是[batch size, num_heads, seq length, head_dim]。计算后结果矩阵大小[batch size, num_heads, seq length, seq length]。算力需求:2 * head_dim * batch size * num_heads * seq length * seq length = 70,368,744,177,664 ≈ 70 TFLOPs。
5.除以head_dim^0.5:
a.计算量:batch size * num_heads * seq length * seq length ≈0.25 TFLOPs。
6.过softmax:
a.实际做了指数计算、加法计算、矩阵标量除法计算;
b.计算量:粗略估计3 * 0.25 TFLOPs。
7.注意力矩阵 * V矩阵:
a.注意力矩阵大小:[batch size, num_heads, seq length, seq length];
b.V矩阵大小:[batch size, num_heads, seq length, head_dim];
c.矩阵乘法结果矩阵大小:[batch size, num_heads, seq length, head_dim];
d.计算量:2 * seq length * batch size * num_heads * seq length * head_dim ≈ 70 TFLOPs。
8.过一层线性层:
a.上一步矩阵大小:[batch size, num_heads, seq length, head_dim];
b.线性层矩阵大小:[hidden_size, num_heads, head_dim];
c.结果矩阵大小:[batch size, seq length, hidden_size];
d.计算量:2 * hidden_size * batch size * seq length * hidden_size ≈ 17.6 TFLOPs。
综上,Attention块的算力需求约为 80层 * (17.6 TFLOPs + 2*2.2 TFLOPs + 70 TFLOPs + 0.25 TFLOPs + 3 * 0.25 TFLOPs + 70 TFLOPs + 17.6 TFLOPs) ≈ 14 PFLOPs。
计算公式可以化简为:num_hidden_layers * batch size * seq length * hidden size * (4.5 * hidden size + 4 * seq length)。
单个FFN块(参数量占比80%,计算量占51%)
# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Qwen2
输入:[batch size, seq length, hidden size]
输出:[batch size, seq length, hidden size]
实际上涉及到计算量的就一行代码,三次矩阵乘法,一次过激活函数,一次矩阵点乘。
拆解分析:
1.输入矩阵 * up_proj矩阵,结果矩阵大小[batch size, seq length, intermediate_size]。计算量2 * batch size * seq length * hidden size * intermediate_size = 63,496,796,504,064 ≈ 63 TFLOPs;
2.输入矩阵 * gate_proj矩阵,同第一步,63 TFLOPs;
3.上一步的结果过激活函数。几乎不需要算力;
4.矩阵点乘,c * batch size * seq length * intermediate_size ≈ 7*c GLOPs;
a.这里c主要是涉及到矩阵点乘的底层实现,我个人估测大概是10以内;
5.上一步的结果矩阵 * down_proj矩阵,同第一步,63 TFLOPs。
综上,FFN层的算力需求约 80层 * (3 * 63 TFLOPs + 7 * c GLOPs) ≈ 15 PFLOPs。
计算公式可以简化为6 * batch size * seq length * hidden size * intermediate_size * num_hidden_layers。
其他杂项(参数量占比0%*80,计算量占0%)
# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Qwen2
输入:[batch size, seq length, hidden size]
输出:[batch size, seq length, hidden size]
输入的正则化和Attention块结果的正则化,本质都是RMSNorm计算。
qwen2的RMSNorm粗略可以理解为做了以下几步:
- 将数据转为双精度。
- 计算整个层所有输出向量的平方平均值向量variance。
计算量:k * batch size * seq length * hidden size ≈ k * 10e9 FLOPs
这里k主要是涉及到平方、均值等运算的底层实现方式,我个人估测大概是10以内。
- 对所有向量乘以向量variance的平方根倒数在加上个eps。
计算量:2 * batch size * seq length * hidden size ≈ 2 * 10e9 FLOPs
- 将数据转回原始精度。
计算量大约是80*2* ( k + 2 ) GFLOPs。代入k=10,粗略估算得这里计算量约 1 TFLOPs。
输出层/分类头/Embedding逆映射(参数量占比1.7%,计算量占1%)
这里先Norm一下,( k + 2 ) GFLOPs
再转输出解码。
隐藏层状态 * [hidden size * vocab_size]
计算量:2 * hidden size * batch size * seq length * vocab_size = 326,554,953,449,472 ≈ 0.3 PFLOPs
公式推导
算力主要是attention、FFN、解码过程产生。
算力合计汇总:
batch size * seq length * hidden size * (2 * vocab_size + num_hidden_layers * (4.5 * hidden size + 4 * seq length + 6 * intermediate_size))
代入到本文的例子,qwen2-72b。4 * 32768 * 8192 * (2 * 152064 + 80 * (4.5 * 8192 + 4 * 32768 + 6 * 29568)) = 29,991,378,670,845,952 ≈ 30 PFLOPs
注:启用GQA后kv cache技术在继续生成阶段约能节约0.3PFLOPs的计算量,1%左右。影响不大。
另,为了简化公式方便理解,将其转变为:
bs*s*h(2V+L(4.5h+4s+6is))
其中,简写含义如下表所示。对于一个已经训好的模型,用户能够干涉的只有batch size和seq length,且按照经验,intermediate_size一般会和hidden size存在一定倍数关系。
进一步化简:bs * s,假设只过一轮epoch,就是全部数据的token数量。
由之前文章分析可得,一个大模型绝大部分参数都在输入输出两个embedding层加transformer层。具体来说,就是attention块和FFN块,和两个词汇映射矩阵。以qwen2为例,加起来共计需要 2*h*V+h*L*(3 is+2.25h)。
代入公式,得正向传播过程总计算量为 2T(2hLs+P) = T(2.6M*s + 144 B),其中T为数据集token数量,P为模型参数量。M表示1e6,B表示1e9
注1:这里公式里的s,就是八股文经常背的,attention原生的平方复杂度影响。也可以看出,只有当s长度超过三位数时,才会对大模型的执行时间产生明显影响。
注2:从中也可以看出,seq长度对大模型的影响并没有那么大。以wen2为例,seq就是拉到32768,对比seql长度为1,总的算力需求也就扩大1.6倍。没那么夸张。
数据验证
官方给出的部署效率(https://huggingface.co/Qwen/Qwen-72B-Chat)如下图所示:
两张A100,BF16,理论算力1248TFLOPS。大模型正向传播需要的算力代入公式:bs*s*(2.6M*s + 144 B),bs取1,s取1000token。则算出来一次正向输出需要0.115秒gpu时长。即输出速度为 8.67qps。和图片中8.48qps基本一致。
拓展八股文:为什么扩大batch size,大模型输出速度先提高后不变
1.batch size过小时,gpu能发挥出的算力和batch size基本正比。
a.结论是:卡内存带宽了。
b.gpu处理数据在一个完整的batch分两个部分,数据转移,数据计算。其中数据转移部分吃内存带宽,需要转移完整的模型参数和一个batch的数据参数。计算部分吃gpu算力,和待计算数据量正相关。
c.小batch下之所以卡在内存带宽,是因为每次转移完整的模型参数这步的处理时间是固定的。而转移batch数据,计算batch数据都和数据量正相关。导致batch越小,转移完整模型参数这步的带宽占比就越大。gpu处理每个batch时,如果batch size过小,数据计算会特别快的完成,但是还需要等待固定时间的模型数据转移完成,导致计算单元空置,算力利用率MFU自适应降低。
d.理论上会有个batch size的切换点。这个点往上,显存占用增加,运算效率不变。但具体怎么算这个点,笔者还没什么思路。
2.batch size达到一定大小,能完全发挥gpu算力时,调用大模型的训练时间已经和batch size无关了。公式里并没有bs。
反向传播过程
结论:反向传播的算力需求一般是前向传播的2倍。
实验证据:
图片来源自笔者之前几篇文章的实验部分。
理论计算:
注:这里其实也可以出一道八股文。为什么反向传播时,正向的矩阵乘法需要2次反向过程的矩阵乘法。
由于大模型算力需求核心在矩阵乘法。因此只考虑涉及到矩阵乘法的反向传播梯度计算。
假设我们有这样的简化模型流程,需要进行一次反向传播:
Y1 = W1 * X
Yo = W2 * Y1
L = loss(Yo,y)
我们通过整个模型的输出Yo计算出了loss L,对loss进行梯度下降。
首先我们需要计算在Yo输出层,承担的对loss梯度。即当前节点/层对最终输出承担多少责任。
也就是计算 △ Yo = d L / d Yo。这个计算根据loss函数的性质来,但计算量不大。(或者说正是因为计算量不大,才会被选为loss函数。)
然后,就可以计算出,W2权重矩阵需要承担多少责任。即,计算个d L / d W2 = △ Yo * Y1。这里是第一次矩阵计算,且每一层都需要算一次。
但这一步还没结束。因为模型是多层模型,需要进一步往前推导。即,需要判断一下W1的责任。而这需要走一步的中介,即先判定Y1的责任。
所以这时需要计算 △ Y1 = d L / d Y1 = W2 * △ Yo。这里是第二次矩阵计算。且除了第一层,每一层都需要计算一次。
这两次矩阵计算和原来的计算矩阵大小相等。所以算力需求一致。由于大模型层数都比较多,所以虽然第二个矩阵计算在第一层不用做,但影响不大。
实际上就是一次矩阵计算(Yo = W2 * Y1)会对应两次反向的矩阵计算(d L / d W2 = △ Yo * Y1)和(△ Y1 = d L / d Y1 = W2 * △ Yo),也就是两倍关系。
梯度更新过程
梯度更新过程需要的计算量,会根据优化器的不同而有所差别。例如随机梯度下降,会对所有的参数,计算一次梯度乘学习率,再计算一次结果加到参数权重。也就是每个参数需要2FLOPs。整个大模型一次更新需要2 * 72b FLOPs。这个数看着很大,但对比前向传播的计算量,也就是忽略不计。
如果是Adam这种,带二阶动量,计算公式比较复杂的,计算量需求又会有所不同。
如上图所示,会经过5次公式。每个公式又有多次计算。但总的来说,一个参数更新一次的计算量需求也是常数级别的。对比前向传播的计算量,也就是忽略不计。
数据验证
官方给出的T数据集大小为7T tokens,seq最长32768。假设大集群的算力利用率MFU为50%。代入公式,得:
总计算量需求:3*7e12*(2.6e6*32,768+144e9) FLOPs
单卡算力:300e12 FLOPS
需要 15.93e9卡秒 ≈ 4,426,000卡小时 ≈ 180000卡天
也就是6k张卡,30天能完成一轮完整训练。
注:训练语料的长度也是在最后阶段才从4096拓展到32768。所以本文预估的算力需求会有一定程度高估,不超过1.6倍。但由于笔者也没找到这个预训练后阶段是什么切分点,所以就按照上限预估了。
而对比llama2时的训练过程(qwen2的模型架构和llama2很相似,可以近似类比)。meta用了1720320卡小时,上下文长度4096,2T数据集(https://llama.meta.com/llama2/)训练一版70B模型。和上面结论基本一致。
读者福利大放送:如果你对大模型感兴趣,想更加深入的学习大模型**,那么这份精心整理的大模型学习资料,绝对能帮你少走弯路、快速入门**
如果你是零基础小白,别担心——大模型入门真的没那么难,你完全可以学得会!
👉 不用你懂任何算法和数学知识,公式推导、复杂原理这些都不用操心;
👉 也不挑电脑配置,普通家用电脑完全能 hold 住,不用额外花钱升级设备;
👉 更不用你提前学 Python 之类的编程语言,零基础照样能上手。
你要做的特别简单:跟着我的讲解走,照着教程里的步骤一步步操作就行。
包括:大模型学习线路汇总、学习阶段,大模型实战案例,大模型学习视频,人工智能、机器学习、大模型书籍PDF。带你从零基础系统性的学好大模型!
现在这份资料免费分享给大家,有需要的小伙伴,直接VX扫描下方二维码就能领取啦😝↓↓↓
为什么要学习大模型?
数据显示,2023 年我国大模型相关人才缺口已突破百万,这一数字直接暴露了人才培养体系的严重滞后与供给不足。而随着人工智能技术的飞速迭代,产业对专业人才的需求将呈爆发式增长,据预测,到 2025 年这一缺口将急剧扩大至 400 万!!
大模型学习路线汇总
整体的学习路线分成L1到L4四个阶段,一步步带你从入门到进阶,从理论到实战,跟着学习路线一步步打卡,小白也能轻松学会!
大模型实战项目&配套源码
光学理论可不够,这套学习资料还包含了丰富的实战案例,让你在实战中检验成果巩固所学知识
大模型学习必看书籍PDF
我精选了一系列大模型技术的书籍和学习文档(电子版),它们由领域内的顶尖专家撰写,内容全面、深入、详尽,为你学习大模型提供坚实的理论基础。
大模型超全面试题汇总
在面试过程中可能遇到的问题,我都给大家汇总好了,能让你们在面试中游刃有余
这些资料真的有用吗?
这份资料由我和鲁为民博士(北京清华大学学士和美国加州理工学院博士)共同整理,现任上海殷泊信息科技CEO,其创立的MoPaaS云平台获Forrester全球’强劲表现者’认证,服务航天科工、国家电网等1000+企业,以第一作者在IEEE Transactions发表论文50+篇,获NASA JPL火星探测系统强化学习专利等35项中美专利。本套AI大模型课程由清华大学-加州理工双料博士、吴文俊人工智能奖得主鲁为民教授领衔研发。
资料内容涵盖了从入门到进阶的各类视频教程和实战项目,无论你是小白还是有些技术基础的技术人员,这份资料都绝对能帮助你提升薪资待遇,转行大模型岗位。
👉获取方式:
😝有需要的小伙伴,可以保存图片到VX扫描下方二维码免费领取【保证100%免费】
相信我,这套大模型系统教程将会是全网最齐全 最适合零基础的!!
更多推荐
所有评论(0)