倒推的顺序建议固定为三步:先把参数量和每参比特数这两列定下来,算出只放权重时的显存下限;再把上下文长度、批大小带来的 KV 缓存加进去;最后在目标环境中实际加载一次,用 nvidia-smi 读到的占用去校正前两步的估算值。三步里最容易出错的是只算权重,往往在首次推理分配 KV 缓存时,显存超出导致 OOM(内存不足)。下面的拆法不绑定具体模型规格,需要你把 MemT 1.0 的实际配置填进占位符,再按结果买卡或调小上下文。
判断方向:先按“权重 + KV 缓存 + 运行时余量”三项分别估算,再把估算结果当成下限而非目标值。适用场景是单机单卡或单机多卡本地加载;操作动作是把参数量、每参比特数、层数、KV 头数、上下文长度、批大小写进统一脚本;验证方式是加载后对比实际占用与估算值的差额;风险边界是估算无法覆盖框架临时拷贝和碎片,最终仍要留出一段空闲显存再决定是否上线。
把参数量与量化档位写成两列
估算的输入变量先只需要两个:总参数量(以十亿为单位,记作 N)和每参比特数(bits per parameter)。权重体积近似等于 N × 每参比特数 ÷ 8 字节,量化档位就是改第二列。同一套公式对不同模型通用,换成别的参数量只需要改 N。
| 量化档位 | 每参比特数(取值区间) | 每 10 亿参数约占用 | 说明 |
|---|---|---|---|
| FP16 / BF16 | 16 | 约 2 GB / 1.86 GiB | 未量化,最直观的基线 |
| INT8 | 8 | 约 1 GB / 0.93 GiB | 部分层可能仍保留高精度 |
| INT4(分组量化) | 4 ~ 4.5 | 约 0.5 ~ 0.56 GB | 分组缩放系数、零点会额外占一点 |
| 混合精度(部分层不量化) | 按层加权 | 介于上述档位之间 | embedding、lm_head、norm 常不量化 |
注意第二列很少是纯整数。INT4 分组量化在每组里额外存 scale 和 zero,实际每参比特数通常略高于 4;模型里未被量化的层会按原精度计入。如果你只能拿到“4bit”这种粗粒度信息,建议先把每参比特数取到区间上沿,宁可估高一点。
按权重占用估算显存下限
权重基线公式为:权重 GiB = N × 1e9 × 每参比特数 ÷ 8 ÷ 1024³。除以 1024³ 得到 GiB,如果显卡按 GB 标稱,要注意两者相差约 7%,记录时统一用一个单位,避免后面核对时对不上。下面是一段可以直接改参数运行的脚本骨架,n_params_b、bits、overhead 都是替换项:
# 通用估算骨架:只算权重,不含 KV 缓存
N_PARAMS_B = 13.0 # 总参数量,单位:十亿,替换成实际值
BITS = 4.5 # 每参比特数,按量化档位填写
OVERHEAD = 1.06 # 加载期额外开销的保守系数,建议 1.05~1.10
weight_gib = N_PARAMS_B * 1e9 * BITS / 8 / (1024 ** 3)
print(f"权重基线: {weight_gib:.2f} GiB")
print(f"含加载开销: {weight_gib * OVERHEAD:.2f} GiB")
误差主要来自三处:一是量化元数据与未量化层让实际比特数偏离名义值;二是加载过程中框架可能先把权重读进高精度张量再转换,出现瞬时高于最终值的占用;三是显存分配器与碎片化,看到的 reserved 往往大于实际分配。所以脚本输出只当作下限,最终留多少要用第四节的实测去校正。
把上下文长度带来的缓存增量计入
KV 缓存大小取决于四个变量:层数、KV 头数(不是注意力头数,分组查询注意力下两者不同)、每头维度、以及上下文长度与批大小。常见估法为:缓存 GiB ≈ 2 × 层数 × KV 头数 × 每头维度 × 上下文长度 × 批大小 × 每元素字节数 ÷ 1024³。系数 2 对应 key 和 value 各一份;每元素字节数按缓存精度取 2(FP16/BF16)或 1(INT8 KV)。把这些变量集中写成可替换配置,方便试不同上下文档位:
# KV 缓存参数:所有值替换成模型配置里的实际字段
CFG = {
"num_layers": 40, # 层数
"num_kv_heads": 8, # KV 头数,GQA 下小于注意力头数
"head_dim": 128, # 每头维度
"ctx_len": 8192, # 上下文长度,调小先试 2048/4096
"batch": 1, # 并发批大小
"bytes_per_elem": 2, # 2=FP16, 1=INT8 KV
}
kv_gib = (2 * CFG["num_layers"] * CFG["num_kv_heads"] * CFG["head_dim"]
* CFG["ctx_len"] * CFG["batch"] * CFG["bytes_per_elem"]) / (1024 ** 3)
print(f"KV 缓存: {kv_gib:.2f} GiB")
上下文长度是线性放大的,批大小也是;把上下文从 8k 降到 4k,缓存大约减半,但权重部分不变,所以长上下文场景下缓存可能反超权重。如果框架支持分页 KV 或 KV 量化,先把对应开关打开再测,否则估算值会偏高而误判为放不下。
实际加载一次并记录占用
校正估算最直接的办法是加载一次并盯着显存。先起一个持续采样的窗口,再在另一个终端启动加载:
# 终端 A:每秒打印一次显存占用
nvidia-smi `--query-gpu`=index,memory.used,memory.total `--format`=csv -l 1
# 终端 B:加载模型,观察 A 的输出变化
# 记录三个时刻:加载前、加载完成后、首次推理返回后
nvidia-smi `--query-gpu`=memory.used `--format`=csv,noheader
记录时至少填下面几项,其中“峰值阶段”一栏要明确写清是加载阶段还是首次推理阶段,这一项决定了你该给权重多留余量还是给缓存多留余量:
| 记录项 | 填写内容 |
|---|---|
| 参数量 / 量化档位 | 例如 13B / INT4 分组量化 |
| 上下文长度 / 批大小 | 实际加载时使用的值 |
| 加载前已占用 | 单位 GiB,排除桌面或其它进程 |
| 权重加载完成后 | 与第二节的权重基线对比差额 |
| 首次推理后峰值 | 与第三节的 KV 缓存估算对比 |
| 峰值出现阶段 | 加载阶段 / 首次推理阶段 |
| 是否启用分页 KV、KV 量化 | 开 / 关,影响对比口径 |
如果实测峰值明显高于估算,先排查是否是加载期的高精度临时拷贝,还是并发请求把批大小推高了;如果实测低于估算,多半是估算时把未量化层算重了。把差额记下来,同一台机器上换上下文长度或批大小时可以直接复用这个差额做修正,不必每次都从零估算。