大模型推理显存和计算量估计方法
最近做吞吐量调试涉及到输入 batch_size 的设置，为了把算力和显存用起来，同时不触发 out of memory，需要提前估计大模型推理过程中的显存占用，我参考了 MindIE 官网的这个文档完成了估计：https://www.hiascend.com/document/detail/zh/mindie/20RC2/mindieservice/servicedev/mindie_service0105.html 。
显存估计
大模型推理的过程中，显存主要用来存储 kvcache。kvcache 的大小和 token 的数量成正比，我们首先来看一下单个 token 的 kvcache 怎么计算。假设单 token 输入经过 embedding 层之后的维度是 hidden_size，那么接下来就要和 k 权重矩阵和 v 权重矩阵相乘，kv 的权重矩阵 shape 是[hidden_size, hidden_size]，所以计算得到的 k 和 v 的维度是 hidden_size。需要注意的是，虽然这里 kv 矩阵只有 1 个，但实际上现在大部分模型都采用了多头注意力，注意力头的数量是 num_attention_heads，单个注意力头的权重矩阵是[hidden_size， hidden_size/num_attention_heads]。此外，如果采用分组 kv 头，也就是模型的 config.json 文件中包含 num_key_value_heads 参数，那么 kv 的权重矩阵 shape 就是[hidden_size, (hidden_size/num_attention_heads)*num_key_value_heads]，这样的话要缓存的 kv 的维度就是(hidden_size/num_attention_heads)*num_key_value_heads。由于大模型一般包含多层 transformer，所以还需要乘以层数。总的来说，单个 token 占用的 kvcache 显存大小计算公式如下：
计算了单个 token 的 kvcache 显存大小后，我们就可以计算整个序列所占的显存大小了：
还可以估计所能支持的最大 batch_size：
在估计最大 batch_size 的时候，我们先根据显卡的最大容量（64G）和模型大小估计了可用于缓存 kvcahe 的空间 mem_for_kvcache，然后估算了可以分配多少个 kvcache_block。接着计算一个 sequence 所需要的 kvcache_block 数量，最后用总的 kvcache_block 数量除以一个 sequence 所需要的 kvcache_block 数量，得到的就是支持的最大 batch 数量。
上面的代码运行结果如下：
为了验证理论分析是否正确，我们用 MindIE 跑一下 qwen2.5-7B 模型：
运行日志截图如下：
可以看到，日志中包含了“kv cache will allocate 1.75GB memory”，和我们上面估算的是一致的。
计算量估计
为了估算推理执行耗时，还需要估计模型推理消耗的计算量，然后结合芯片算力估计时延。
Transformer 模型的计算量主要来自自注意力机制、前馈网络（FFN）和最后的 lm_head 层。假设模型参数如下：
L：Transformer 层数。
H：隐藏层大小（隐藏维度）。
I：FFN 中间层大小（通常 I>H）。
V：词表大小。
S：prefill 输入序列长度。
T：序列总长度（prefill+生成 token）。
Prefill 计算量
计算每层 Transformer 的 FLOPs：
首先计算自注意力层。第一步是求 q、k、v，计算公式是 input*weight，input 的 shape 是[S, H]，weight 的 shape 是[H, H]，所以 q/k/v 要做 SH 次向量内积，每次向量内积要做 H 次乘法和（H-1）次加法，近似于 2H，所以 q、k、v 所有的计算量是 3*SH*2H=6SH^2。需要注意的是，如果采用的是分组 kvcache，那么计算 q、k、v 的时候，H 要换成(H/num_heads)*num_kv_heads，但是在估算的时候可以近似为 H；然后是 Q*K（实际上是 Q 乘以 K 的转置），输入 shape 是[S, H]和[H, S]，计算量是 2S^2H；接着是 Q*K*V，输入 shape 是[S, S]和[S, H]，计算量是 2S^2H；最后是输出投影，输入 shape 是[S, H]和[H, H]，计算量是 2SH^2；所以自注意力层的计算量是 8SH^2+4S^2H。
然后计算 FFN 层的计算量，FFN 层包含一个升维层、一个 gate 层和一个降维层，计算量分别为 2SHI、2SHI、2SIH，所以总计算量为 6SHI。
所以每个 transformer 层的计算量为 8SH^2+4S^2H+6SHI。再加上最后的 lm_head，prefill 的计算量为 L(8SH^2+4S^2H+6SHI)+2SHV。
Decode 计算量
和 prefill 相比，decode 的主要变化是输入序列长度为 1。
首先计算自注意力层。计算 qkv 的计算量是 6H^2。计算 q*k 的时候，由于要把缓存的 k 也加上，假设当前序列长度是 T，那么输入 shape 是[1, H]和[H, T]，所以计算量是 2HT；接着计算 qkv，输入 shape 是[1,T]和[T, H]，计算量是 2HT。所以自注意力层的计算量是 8H^2+4HT。
然后计算 FFN 层的计算量，参考 prefill 的过程，可知是 6HI。
所以每个 transformer 层的计算量为 8H^2+4HT+6HI。再加上最后的 lm_head，decode 的计算量为 L(8H^2+4HT+6HI)+2HV。需要注意的是，这里面的 T 是变化的，如果要计算 K 个 decode 过程的平均耗时，可以取 T=(S+S+K)/2 进行估计。
实验验证
为了验证我们的理论公式，我们继续基于 qwen2.5-7B 进行验证，还是运行：
输出结果部分截图如下：
在输入/输出取 1024/1024，batch_size=16 的情况下，首 token 时延 1088.62ms，decode 平均时延 18.77ms。
我们把理论计算公式用 python 实现：
注意，代码中的 npu_cal_ability 是我们使用的 npu 单卡算力，300T Flops 左右。
输出结果如下：
可以看到，prefill 的预估时间为 862ms，和实测的 1088.62ms 比较接近，但是增量平均时延为 0.85ms，远小于实测的 18.77ms。主要有以下原因：
1，AI 芯片的算力参数为峰值算力，一般情况下 AI 芯片的算力利用率在 60%左右；
2，AI 芯片的算力大部分来自矩阵乘法单元，全量计算都是大矩阵运算 GEMM，可以充分利用 AI Core 的能力，但是增量计算都是小矩阵运算（特别是 batch_size=1 的时候，退化为向量矩阵运算 GEMV），导致算力利用率很低；
3，token 时延除了包含计算时间，还有内存搬运时间、软件栈之间的数据传输时间等等，对于 decode 这种运算时间短的场景，其他环节的时延会占很大的比例。
此外，对于分布式推理场景，不能仅凭计算量来估计时延，因为通信算子的耗时往往会较大，比如 MOE 结构模型中的 alltoall 通信在推理过程中可能占总时延的 30%。
还未添加个人签名 2020-11-13 加入
还未添加个人简介
评论