登录/ 注册
 找回密码
 立即注册
关闭

CSDN热搜

程序园 精品问答 技术交流 资源下载
返回列表 发新帖
首页 › 业界区 › 科技 ›大模型推理显存和计算量估计方法

大模型推理显存和计算量估计方法

碛物 2025-6-22 22:35:11
最近做吞吐量调试涉及到输入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显存大小计算公式如下：
  1. import numpy as np
  2. # take qwen2.5-7B as example
  3. hidden_size = 3584
  4. num_attention_heads = 28
  5. num_hidden_layers = 28
  6. num_key_value_heads = 4  # if not grouped kv cache, num_key_value_heads equals num_attention_heads
  7. kvcache_dtype_byte = 2  # bf16/fp16 -> 2, int8 -> 1
  8. cache_num = 2  # k cache and v cache
  10. one_token_cache = hidden_size / num_attention_heads * num_key_value_heads * num_hidden_layers * kvcache_dtype_byte * cache_num
  11. print(f"one_token_cache is: {one_token_cache} byte.")
复制代码
计算了单个token的kvcache显存大小后，我们就可以计算整个序列所占的显存大小了：
  1. input_length = 1024
  2. output_length = 1024
  3. batch_size = 16
  4. allocate_mem = one_token_cache * (input_length + output_length) * batch_size / (1024*1024*1024) # GB
  5. print(f"allocate memory should be larger than {allocate_mem} G.")
复制代码
还可以估计所能支持的最大batch_size：
  1. # estimate max batch_size for qwen2.5-7B
  2. model_size = 7  # B
  3. total_mem = 64  # 64G for each npu
  4. mem_coefficient = 0.8
  5. mem_for_kvcache = (total_mem - model_size*kvcache_dtype_byte)*mem_coefficient*1024*1024*1024  # byte
  6. kvcache_block_size = 128
  7. max_block_num = np.floor(mem_for_kvcache/(kvcache_block_size*one_token_cache))
  8. print("max_block_num: ", max_block_num)
  9. one_sequence_block_num = np.ceil((input_length + output_length) / kvcache_block_size)
  10. max_batch_Size = np.floor(max_block_num / one_sequence_block_num)
  11. print("max_batch_Size: ", max_batch_Size)
复制代码
在估计最大batch_size的时候，我们先根据显卡的最大容量（64G）和模型大小估计了可用于缓存kvcahe的空间mem_for_kvcache，然后估算了可以分配多少个kvcache_block。接着计算一个sequence所需要的kvcache_block数量，最后用总的kvcache_block数量除以一个sequence所需要的kvcache_block数量，得到的就是支持的最大batch数量。
上面的代码运行结果如下：
  1. one_token_cache is: 57344.0 byte.
  2. allocate memory should be larger than 1.75 G.
  3. max_block_num:  5851.0
  4. max_batch_Size:  365.0
复制代码
为了验证理论分析是否正确，我们用MindIE跑一下qwen2.5-7B模型：
  1. bash run.sh pa_bf16 performance [[1024,1024],[256,256]] 16 qwen /home/jinxiulang/qwen2.5/Qwen2.5_7B_Instruct 1
复制代码
运行日志截图如下：
可以看到，日志中包含了“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]，计算量是2SH2；所以自注意力层的计算量是8SH2+4S^2H。
​        然后计算FFN层的计算量，FFN层包含一个升维层、一个gate层和一个降维层，计算量分别为2SHI、2SHI、2SIH，所以总计算量为6SHI。
​        所以每个transformer层的计算量为8SH2+4S2H+6SHI。再加上最后的lm_head，prefill的计算量为L(8SH2+4S2H+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层的计算量为8H2+4HT+6HI。再加上最后的lm_head，decode的计算量为L(8H2+4HT+6HI)+2HV。需要注意的是，这里面的T是变化的，如果要计算K个decode过程的平均耗时，可以取T=(S+S+K)/2进行估计。
实验验证

为了验证我们的理论公式，我们继续基于qwen2.5-7B进行验证，还是运行：
  1. bash run.sh pa_bf16 performance [[1024,1024],[256,256]] 16 qwen /home/jinxiulang/qwen2.5/Qwen2.5_7B_Instruct 1
复制代码
输出结果部分截图如下：
在输入/输出取1024/1024，batch_size=16的情况下，首token时延1088.62ms，decode平均时延18.77ms。
我们把理论计算公式用python实现：
  1. # estimate calculation amount for qwen2.5-7B
  2. # prefill
  3. L = num_hidden_layers
  4. S = input_length
  5. H= hidden_size
  6. I = 18944
  7. V = 152064
  8. attention_cal_prefill = L*(8*S*H*H + 4*S*S*H)
  9. forward_cal_prefill = L*6*S*H*I
  10. head_cal_prefill = 2*S*H*V
  11. total_cal_prefill = (attention_cal_prefill + forward_cal_prefill + head_cal_prefill)*batch_size
  12. print(f"Total calculation amount for prefill is {total_cal_prefill:.3e}, the percentages of attention layer,"
  13.       f" forward layer and head layer are: {attention_cal_prefill/total_cal_prefill}, {forward_cal_prefill/total_cal_prefill}, {head_cal_prefill/total_cal_prefill}")
  15. # decode
  16. T = input_length + output_length / 2
  17. attention_cal_decode = L*(8*H*H + 4*H*T)
  18. forward_cal_decode = L*6*H*I
  19. head_cal_decode = 2*H*V
  20. total_cal_decode = (attention_cal_decode + forward_cal_decode + head_cal_decode)*batch_size
  21. print(f"Total calculation amount for averaged decode is {total_cal_decode:.3e}, the percentages of attention layer,"
  22.       f" forward layer and head layer are: {attention_cal_decode/total_cal_decode}, {forward_cal_decode/total_cal_decode}, {head_cal_decode/total_cal_decode}")
  24. npu_cal_ability = 300*1e+12
  26. print(f"The calculation time for prefill and decode are: {(total_cal_prefill / npu_cal_ability)}, {total_cal_decode / npu_cal_ability}")
复制代码
注意，代码中的npu_cal_ability是我们使用的npu单卡算力，300T Flops左右。
输出结果如下：
  1. Total calculation amount for prefill is 2.586e+14, the percentages of attention layer, forward layer and head layer are: 0.20832364566379913, 0.7226226458963032, 0.06905370843989769
  2. Total calculation amount for averaged decode is 2.558e+11, the percentages of attention layer, forward layer and head layer are: 0.21849896717925177, 0.7133348634381456, 0.06816616938260271
  3. The calculation time for prefill and decode are: 0.8620572025378134, 0.000852813851306666
复制代码
可以看到，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%。
本文由博客一文多发平台 OpenWrite 发布！

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

使用道具 举报
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册
发帖
碛物
2025-6-22 22:35:11

0

粉丝关注

21

主题发布

板块介绍填写区域，请于后台编辑
微信扫一扫