大模型训练计算量及训练时间估计

date

Apr 1, 2025

slug

transformer_computation_and_training_cost

TODO

后向的计算量大约是前向的两倍

🥰总结

⁉️问题

🧐内容

缩写

缩写	符号	描述
b	batch size	批大小
s	sequence length	token的长度
v	vocab size	词表大小
h	hidden size	隐层大小,
d	head dim	head的维度
a	number heads	MHA中head的数量
l	number layers	层数

计算量

这里把乘加分开为两次运算, 一个矩阵乘计算量为2mnk

self-attetion

计算：矩阵乘法的输入和输出形状为。计算量为。

矩阵乘法的输入和输出形状为。计算量为。

计算在 V 上的加权，矩阵乘法的输入和输出形状为。计算量为。

attention后的线性映射，矩阵乘法的输入和输出形状为。计算量为。

MLP

第一个线性层，矩阵乘法的输入和输出形状为。计算量为。

第二个线性层，矩阵乘法的输入和输出形状为。计算量为。

其他

另一个计算量的大头是logits的计算，将隐藏向量映射为词表大小。矩阵乘法的输入和输出形状为，计算量为

汇总

因此，对于一个层的transformer模型，输入数据形状为 [b,s] 的情况下，一次训练迭代的计算量为:

计算量与参数量

当隐藏维度 h 比较大，且远大于序列长度 s 时，我们可以忽略一次项，计算量近似为

模型参数量提到模型参数量近似为，输入的tokens数为 bs ，存在等式:

我们可以近似认为：在一次前向传递中，对于每个token，每个模型参数，需要进行2次浮点数运算，即一次乘法运算和一次加法运算。

一次训练迭代包含了前向传递和后向传递，后向的计算量大约是前向的两倍。因此，前向传递 + 后向传递的系数 =1+2=3 。一次训练迭代中，对于每个token，每个模型参数，需要进行 2∗3=6 次浮点数运算。

接下来，我们可以估计训练GPT3-175B所需要的计算量。对于GPT3，每个token，每个参数进行了6次浮点数运算，再乘以参数量和总tokens数就得到了总的计算量。GPT3的模型参数量为 174600M ，训练数据量为 300B tokens。

后向的计算量大约是前向的两倍

https://zhuanlan.zhihu.com/p/675517271

训练时间估计

模型参数量和训练总tokens数决定了训练transformer模型需要的计算量。给定硬件GPU类型的情况下，可以估计所需要的训练时间。给定计算量，训练时间（也就是GPU算完这么多flops的计算时间）不仅跟GPU类型有关，还与GPU利用率有关。计算端到端训练的GPU利用率时，不仅要考虑前向传递和后向传递的计算时间，还要考虑CPU加载数据、优化器更新、多卡通信和记录日志的时间。

一般来讲，GPU利用率一般在 0.3∼0.55 之间。

上文讲到一次前向传递中，对于每个token，每个模型参数，进行2次浮点数计算。使用激活重计算技术来减少中间激活显存需要进行一次额外的前向传递，因此前向传递 + 后向传递 + 激活重计算的系数=1+2+1=4。使用激活重计算的一次训练迭代中，对于每个token，每个模型参数，需要进行 2∗4=8 次浮点数运算。在给定训练tokens数、硬件环境配置的情况下，训练transformer模型的计算时间为：

以GPT3-175B为例，

参数量	GPU	显存	GPU数量	峰值算力	tokens数
175B	A100	40GB	1024	312TFLOPS	300B

设GPU利用率为0.45，则所需要的训练时间为:

以LLaMA-65B为例，

参数量	GPU	显存	GPU数量	峰值算力	tokens数
65B	A100	80GB	2048	624TFLOPS	1.4TB

设GPU利用率为0.3，则所需要的训练时间: