大模型学习

大模型的参数量和显存占用估算

现在业界的大语言模型都是基于transformer模型的，模型结构主要有两大类：encoder-decoder（代表模型是T5）和decoder-only，具体的，decoder-only结构又可以分为Causal LM（代表模型是GPT系列）和Prefix LM（代表模型是GLM）。针对decoder-only框架，估算其参数量和显存占用。

参数量约为 $12lh^{2}$ ，其中l指transformer层数，h指隐藏层维度。

训练显存占用约为20*参数量，单位B。20=2+4+2+4+4+4，前两个数字是权重，接着两个是梯度，最后两个是优化器状态大小。每个可训练模型参数都会对应1个梯度，并对应2个优化器状态。在混合精度训练中，会使用float16的模型参数进行前向传递和后向传递，计算得到float16的梯度；在优化器更新模型参数时，会使用float32的优化器状态、float32的梯度、float32的模型参数来更新模型参数。

推理显存占用约为2*参数量，单位B。如果使用KV cache来加速推理过程，KV cache也需要占用显存，约为 $4blh(s+n)$ ，b是batch，l是transformer层数，h指隐藏层维度，s是输入序列长度，n是输出序列长度，4=2*2，k和v的cache，每个cache fp16存储，占用2B。

计算量FLOPs约为

计算量和参数量的关系，近似认为，在一次前向传递中，对于每个token，每个模型参数，需要进行2次浮点数运算，即一次乘法法运算和一次加法运算。一次训练迭代包含了前向传递和后向传递，后向传递的计算量是前向传递的2倍。因此，一次训练迭代中，对于每个token，每个模型参数，需要进行6次浮点数运算。

训练时间估计参考下面的公式，8是因为激活重计算技术来减少中间激活显存需要进行一次额外的前向传递，即4*2次浮点数运算。

中间激活的显存占用大小约为 $(34bsh+5bs^{2}a)*l$ ，其中b是batch，s是序列长度，a是注意力头数，l是transformer层数，h指隐藏层维度。在训练过程中中，模型参数（或梯度）占用的显存大小只与模型参数量和参数数据类型有关，与输入数据的大小是没有关系的。优化器状态占用的显存大小也是一样，与优化器类型有关，与模型参数量有关，但与输入数据的大小无关。而中间激活值与输入数据的大小（批次大小和序列长度是成正相关的，随着批次大小和序列长度的增大，中间激活占用的显存会同步增大。