显存计算

显存计算

推理显存计算公式

推理模型显存占用 = 模型权重 （Model Weight 固定占用）+ KV Cache （随上下文增长） + 其他开销（Overhead 约10%）

模型权重计算

一个全精度float32类型的参数占用4个字节

1B模型有10亿的参数，每个参数占4字节

1G=1024M=$1024^2$K = $1024^3$ Byte

$10^9$ x 4 ÷ $1024^3$ = 3.725G (约4G)

全精度部署1B模型占用4G显存

btype	1B（10亿）参数需要占用内存	7B	13B
Float32（全精度）	4G	28G	52G
fp16/bf16（半精度）	2G	14G	26G
int8	1G	7G	13G
int4	0.5	3.5G	6.5G

KV Cache显存

KV Cach缓存作用，减少重复的推理，提高推理的效率

KV Cache所占有的显存 = 每个Token生成时需要的显存 x 生成Token的数量 x 并发用户量

每个Token所需的显存= 模型层数 x hidden dimension（隐藏维度） x 字节数（全精度就是4Byte）x 2 （KV Cache 由 Key 和 Value 两个部分组成，所以要乘以2）

生成Token数 = 生成上下文数（输入 + 输出字长）

例子：每个Token所需的显存 = 模型层数（80）x hidden dimension (隐藏维度)（8192） x 2Byte x 2 = 2.5MB

KV Cache 显存 ≈ 2 × 层数 × 隐藏层维度 × 上下文长度 × 2 ÷ (1024³) GB

每个用户都需要开辟新的KV Cache

KV Cache 所需的显存 = 每个Token所需的显存为 2.5MB x 32k上下文 x 10个用户 = 800GB

隐藏层维度

注 模型可能会有多种推理模式，每种模式的隐藏层维度都不一样，需要根据使用情况具体进行判断

隐藏层维度 = KV注意力头数 (num_key_value_heads) x 单头维度(head_dim )

隐藏层维度数可以从模型的config.json里的hidden_size里看到

层数

num_hidden_layers: 40

其他开销

activation：神经网络在每层计算时的激活函数所占空间

buffers 缓冲区，临时存放中间变量

overheads 碎片开销，显存空间碎片

activation + buffers + overheads 所需显存 约等于 （Model Weight + KV Cache）* 10%

训练时显存计算

模型参数

梯度值

优化器状态

激活值

训练数据

缓存和驱动

量化和蒸馏

量化和蒸馏本质都是压缩

量化

量化（Quantization）将一个参数值从高精度转为低精度的过程，例如一个浮点数从float32（4Byte）转为float16（2Byte），以牺牲准确率来换取时间和空间

蒸馏

蒸馏是小模型（Student Model）模仿 大模型（Teacher Model），通过大模型的问题和回复去训练小参数的模型，使得小参数模型能够和大模型在一定程度上相似。

模型并行部署的方式