分析transformer模型的参数量、计算量、中间激活、KV cache、bf16、fp16、混合精度训练

转载已于 2025-05-11 00:55:39 修改 · 1.5w 阅读

34 ·

CC 4.0 BY-SA版权

原文链接：https://zhuanlan.zhihu.com/p/624740065

文章标签：

#transformer #深度学习 #人工智能

于 2023-07-30 19:49:30 首次发布

难得一遇的好文，转载自https://zhuanlan.zhihu.com/p/624740065，先做个简单总结：

训练时的参数量由模型参数、后向的梯度、优化器参数、前向的中间激活组成：

假设采用Zero3，模型参数量为B，混合精度训练，卡数为N：

优化器参数、后向的梯度、模型参数是DeepSpeed中分的参数，每块卡分20B/N，关于20B怎么来的可以参考 Deepspeed、ZeRO、FSDP、ZeRO-Offload、all reduce、reduce-scatter。简单记住20B/N或者16B/N即可，假设8卡A100训练一个1B的模型，这部分每张卡占用4G*20/8=10G左右，你的A100里的10个G就没有了。
前向的中间激活才是真的大头。中间激活占用的显存大小可以近似为：
$5bs^2a) \cdot l$ ，其中b是batch_size，s是序列长度，l是模型层数，a是有几个头，如果transformer模型的参数量近似为 $12lh^2$ ，可以直观感受一下比例。算的时候需要计算器没法口算，但是要知道中间激活是和batch_size 100%成正比，和模型头数是一部分成正比，和序列长度一部分正比，另外一个部分平方成正比。以GPT3 2048序列长度感受一下，bs=1是中间激活是模型大小的0.79倍。而且中间激活是每个卡独占算的时候不能像上面Zero3一样给除N。这就是为啥8卡Finetune一个7B的模型，序列长度是4096，如果bs=1时，可能每张卡大概中间激活大小就超过模型的2倍，也就是至少占了50多G的显存，bs进一步大的时候，就直接OOM了。

模型参数

假设Transformer的hidden_size是h，那么总的一层transformer层的参数量为 $12h^2+13h$ ，l层的transformer计算量就是 $l*(12h^2+13h)$ :

多头注意力四个参数矩阵Q、K、V、O和他们的bias对应 $4h^2+4h$
position wise feedforward对应 $4h^2+4h+4h^2+h$
多头注意力和position wise feedforward各有一个LayerNorm对应2个可训练模型缩放参数beta和平移参数gamma参数是4h

前向计算过程中产生的中间激活

前向计算过程中产生的中间激活，中间激活值与输入数据的大小（批次大小b和序列长度 l）是成正相关的，随着批次大小b和序列长度l的增大，中间激活占用的显存会同步增大。当我们训练神经网络遇到显存不足OOM（Out Of Memory）问题时，通常会尝试减小批次大小来避免显存不足的问题，这种方式减少的其实是中间激活占用的显存。

前向的中间激活才是真的大头。中间激活占用的显存大小可以近似为：
$5bs^2a) \cdot l$ ，其中b是batch_size，s是序列长度，l是模型层数，a是有几个头，如果transformer模型的参数量近似为 $12lh^2$ ，可以直观感受一下比例。算的时候需要计算器没法口算，但是要知道中间激活是和batch_size 100%成正比，和模型层数100%成正比，和模型头数是一部分成正比，和序列长度一部分正比，另外一个部分平方成正比。

可以看到随着批次大小 $b$ 的增大，中间激活占用的显存远远超过了模型参数显存。通常会采用 激活重计算技术来减少中间激活，理论上可以将中间激活显存从 $O (n)$ 减少到 $O(\sqrt{n})$ ，代价是增加了一次额外前向计算的时间，本质上是“时间换空间”。

后向传递计算得到的梯度、优化器状态

在一次训练迭代中，每个可训练模型参数都会对应1个梯度，并对应2个优化器状态（Adam优化器梯度的一阶动量和二阶动量）。设模型参数量为 $\Phi$ ，那么梯度的元素数量为 $\Phi$ ，AdamW优化器的元素数量为 $2\Phi$ 。float16数据类型的元素占2个bytes，float32数据类型的元素占4个bytes。在混合精度训练中，会使用float16的模型参数进行前向传递和后向传递，计算得到float16的梯度；在优化器更新模型参数时，会使用float32的优化器状态、float32的梯度、float32的模型参数来更新模型参数。因此，对于每个可训练模型参数，占用了 $(2 + 4) + (2 + 4) + (4 + 4) = 20 b y t es$ 。使用AdamW优化器和混合精度训练来训练参数量为 $\Phi$ 的大模型，模型参数、梯度和优化器状态占用的显存大小为 $20\Phi$

推理时的参数量少了梯度、优化器状态、中间激活，但多了kvcache

假设输入序列的长度为s，输出序列的长度为n，以float16来保存KV cache，那么KV cache的峰值显存占用大小为 $b * (s + n) h * l * 2 * 2 = 4 l h (s + n)$ 。这里第一个2表示K/V cache，第个2表示float16占2个bytes。
以GPT3为例，对比KV cache与模型参数占用显存的大小。GPT3模型占用显存大小为350GB。假设批次大小b=64，输入席列长度 =512，输出序列长度n =32，则KV cache占用显存大约为 $46 l h (s + n)$ = 164,282,499,072bytes约等于164GB，大约是模型参数显存的0.5倍

训练时的计算量

FLOPs，floating point operations，表示浮点数运算次数，衡量了计算量的大小

对于 $\in R^{m \times n}$ , $\in R^{n \times p}$ ，矩阵A✖️B浮点运算需要mnp次乘法和mnp次加法，因此FLOPs为2mnp。这个刚好是计算复杂度*2的关系，因此我们下面只估计乘法的计算次数，最后算FlOPs乘2就好了

模型计算复杂度

假设Transformer的hidden_size是h，序列长度是s，batch大小是b，那么总的一层transformer层的参数量为 $12bsh^2+2bs^2h$ ，l层的transformer计算量就是 $l*(12bsh^2+2bs^2h)$ :

多头注意力四个参数矩阵Q、K、V、O对应 $4bsh^2$ （忽略他们的bias）
QK、KV的复杂度是 $2bs^2h$ （ $[b, s, h] * [b, h, s] = [b, s, s]$ 和 $[b, s, s] * [b, s, h] = [b, s, h]$ ，这俩都是 $bs^2h$ ）
position wise feedforward对应 $4bsh^2+4bsh^2$

计算量与参数量的关联

当隐藏维度h比较大，且远大于序列长度s时，我们可以忽略一次项，计算量可以近似为 $\frac{(12bsh^2+2bs^2h)*2}{12h^2+13h}\approx2$ 。我们可以近似认为：在一次前向传递中，对于每个token，每个模型参数，需要进行2次浮点数运算，即一次乘法法运算和一次加法运算。

一次训练迭代包含了前向传递和后向传递，后向传递的计算量是前向传递的2倍（loss.backward()和optimizer.step()两步）。因此，前向传递 + 后向传递的系数 = 1+2 =3 。一次训练迭代中，对于每个token，每个模型参数，需要进行 2*3=6次浮点数运算。

训练时间估计

模型参数量和训练总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模型的计算时间为：
在这里插入图片描述

fp16和bf16的相关问题

为什么有了fp16又整了一个bf16

BF16 的设计思想是在不改变内存占用的情况下，用1/10的精度换取了 10^34倍的值域。
BF16 只有 2bytes 的内存，但其数值范围与 4bytes 的 FP32 相同。在深度学习领域，数值范围的作用远高于数值精度；即数据类型的指数位的作用大于尾数位的作用。采用SGD的网络权重的更新方式为 $\theta_t=\theta_{t-1}-\eta*g_t$ 。由于梯度 $g_t$ 和学习率 $\eta$ 通常较小，往往会出现很小的数值，因此必须使用能够表达较大范围的数据类型。使用 FP16时往往会出现 underflow（下溢）的情况，即小于 − 6.55 × 10^4 的数值被截断为0 ，导致梯度无法更新。所以 BF16 具有比 FP16 更优异的性能。
BF16 与 FP32 的转换很容易。使用混合精度计算时，需要频繁得对 BF16/FP32 和 FP32 进行转换。BF16 基本上可以看作成一个“截断”版的 FP32, 两者之间的转换是非常直接，其实现电路也会非常简单。相比于 FP16，BF16的使用能有效的降低电路的面积。

原文链接：https://blog.youkuaiyun.com/qq_43799400/article/details/134182459

如果还是v卡，继续用fp16做混合精度训练，大致的流程如何？

MIXED PRECISION TRAINING 来自于论文https://arxiv.org/pdf/1710.03740.pdf，
整体流程：FP32权重 -> FP16权重 -> FP16计算前向 -> FP32的loss，扩大 -> 转为FP16 -> FP16反向计算梯度 -> 缩放为FP32的梯度更新权重
在这里插入图片描述
值得注意的是：

Loss scaling：由于梯度值都很小，用FP16会下溢，因此先用FP32存储loss并放大，使得梯度也得到放大，可以用FP16存储，更新时变成FP32再缩放。如果梯度过大导致FP16上溢，Nvidia Apex，或者Pytorch1.6 AMP都会跳过step()然后调整Loss的放大倍数。
在涉及到累加操作时，比如BatchNorm、Softmax，需要用FP32保存，一般使用GPU中TensorCore的FP16*FP16+FP32=FP32运算。这里是用FP32进行存储的主要原因应该是为了避免精度损失，因为如果直接在FP16上进行累加，可能会出现舍入误差（毕竟FP16在每一个数值区间能够表示的最小间隔比FP32的大）。

原文链接：https://zhuanlan.zhihu.com/p/110278004

inference的时候哪些算子不能转换fp16？

当然还是刚才说到的涉及到累加操作时，比如BatchNorm、Softmax，尤其是BatchNorm，这里引用来自apex的一个issue（https://github.com/NVIDIA/apex/issues/122），当然这这个issue也是说的训练时候怎么转，【Quantization aware training(神经网络量化训练)】https://www.bilibili.com/video/BV13s4y1D73L?vd_source=e260233b721e72ff23328d5f4188b304 里面强调做一些量化的时候要把BN算子给先合并掉

以下是转载截屏：
在这里插入图片描述