大模型显存如何估算？

一般会分成两种情况：推理和训练。推理时，我们主要考虑模型权重、KV 缓存和激活显存。而训练相比推理，还需要额外存储优化器状态和梯度，显存需求会更高。

1. 推理

推理的显存主要有三部分：模型权重、KV缓存和激活显存。首先加载模型权重是显存的主要开销。

• 模型权重

  模型大小=参数数量×精度

  比如全精度（FP32）每个参数占 4 字节；半精度（FP16）是 2 字节；更小的量化，比如 int8 和 int4，则分别是 1 字节和 0.5 字节。

  要减少这部分显存，可以用 权重量化 技术，比如将 FP32 的权重压缩成 int8 或 int4。

• KV 缓存

  在解码时每生成一个 token，模型都会用到之前所有的 KV值。为了避免重复计算，需要把这些历史值缓存起来。

  KV 缓存=2×批次大小×序列长度×层数×隐藏层大小×精度

  这里乘2是因为要算一次 K 还要算一次 V

  优化这部分显存，可以用 PagedAttention，通过更高效的机制减少缓存需求。

• 激活显存

  参考：中间激活值显存分析

  在模型的前向传播中，每层的中间激活值都需要存储。直接用Megtron论文里面给的计算公式：