qinghuan
← 回到知识探索

AI 直觉 / 动手实验

权重装得下，为什么一聊天就显存不足？

用一张简化预算拆开权重与 KV 缓存，再看训练状态和临时内存为何也要留空间。

约 9 分钟资料核对：2026-09-26
展开本文目录
  1. 固定的权重，增长的上下文缓存
  2. 只改序列数和长度，观察两项小计
  3. 实际峰值还包括哪些部分
  4. 记住这三点
  5. 别把例子用过头
  6. 自己试一小步
  7. 查原始资料
  8. 找一份资料继续练
开始前，会这些就够了
  • 理解字节与数组大小
  • 知道模型按 Token 逐步生成文本
遇到不熟悉的词，先看这 3 个解释
参数与权重
模型中参与计算、通常由训练学得的数值。权重是其中一类，偏置也属于参数；它们共同决定输入怎样变成输出，而不是一条条可直接查询的知识记录。
上下文
模型在当前计算中能够利用的输入信息。文本生成时，通常包括系统要求、选入的历史消息、检索材料和已生成文本；能同时容纳的容量通常按词元数限制。
KV 缓存
自回归模型逐个生成词元时，保存之前算过的键和值，以便下一步重复使用。键和值是注意力计算中的两组表示；缓存减少重复计算，同时占用额外内存。
打开完整术语小抄 ↗

把模型载入显卡时没有报错，输入一段长材料后却提示显存不足，并不矛盾。权重只是运行时的一部分。处理输入、保存已读上下文的中间结果、生成新内容，都可能继续占空间。因此“模型文件小于显存容量”只能提供一个线索，不能保证指定任务一定跑得动。

先把场景说清：推理是用已学参数完成预测，训练还要计算如何更新参数。两者需要保留的东西不同。本篇先用普通解码器的简化推理预算看清变化，再说明不能直接把这张小表当成训练或所有模型的总显存公式。

固定的权重，增长的上下文缓存

权重可先按参数数量乘每个参数所占字节估算，但量化辅助信息、额外副本和不同层的数据类型仍需另算。生成时常用 KV cache，保存先前位置的键和值，供后续注意力计算复用。这能减少重复计算，同时需要空间；它不是把整段聊天永久写进模型权重。

一个 KV 头，可以理解为注意力中的一组键和值的表示。对于各层形状相同、完整保留上下文的普通 KV 缓存，可以数出元素总量：两份键和值，乘层数、批量中的序列数、缓存长度、KV 头数和每头维度，最后乘每个元素的字节数。这里用的是实际保存 K/V 的头数，不应不加检查就拿查询头数代替。查询头与 KV 头的数量可能不同；其他模型还可能采用不同缓存结构。

只改序列数和长度，观察两项小计

下面完全自拟一套预算数字：一亿个两字节权重，十二层，每层四个 KV 头，每头六十四维。假设权重与缓存都在同一设备，批量共享一份权重，各序列缓存相互独立且长度相同。它不是某个真实型号的配置，也不会创建大数组或占用 GPU；程序只计算两项字节数。

第一种情况是一条序列缓存一千零二十四个位置，第二种是四条序列各缓存四千零九十六个位置。缓存长度包括本次已处理的输入与生成位置；输出只是权重加缓存的小计。MiB 按二的二十次方字节换算，避免与十进制 MB 混用。两项小计不包含其他运行空间，不能直接拿它承诺模型能装进某张卡。

小实验 · python
MiB = 2 ** 20
parameters = 100_000_000
weight_bytes = parameters * 2
layers, kv_heads, head_dim = 12, 4, 64
cache_element_bytes = 2
for batch, cached_tokens in [(1, 1024), (4, 4096)]:
    cache_bytes = (2 * layers * batch * cached_tokens
                   * kv_heads * head_dim * cache_element_bytes)
    print(f"batch={batch} cached={cached_tokens} "
          f"weights={weight_bytes / MiB:.1f} MiB "
          f"cache={cache_bytes / MiB:.1f} MiB "
          f"subtotal={(weight_bytes + cache_bytes) / MiB:.1f} MiB")

怎样运行：Python 3.10+；保存为 memory_budget.py，在该目录运行 python memory_budget.py。只用标准库，无需联网。

运行后对照结果
batch=1 cached=1024 weights=190.7 MiB cache=12.0 MiB subtotal=202.7 MiB
batch=4 cached=4096 weights=190.7 MiB cache=192.0 MiB subtotal=382.7 MiB

实际峰值还包括哪些部分

推理还会产生激活值，也就是各层算出的中间结果，以及算子所需临时空间。训练通常另有梯度、优化器状态和为反向传播保留的激活；它们随算法和精度设置变化，不能统一套“每参数固定多少字节”。只量化权重，也不代表 KV 缓存与这些中间结果同步缩小。

缓存实现同样会改变预算。静态缓存可能提前预留最大长度；滑动窗口只保留特定范围；量化缓存、共享前缀或移到 CPU 又有各自条件。上面的逐元素公式适用于写明的简单布局，未考虑这些策略。换了架构或运行时，应先核对实际张量与分配方式，再讨论能节省多少。

测量时也要分清口径。PyTorch 文档区分张量实际占用与缓存分配器保留的空间，显卡监控值未必等于两者中的任意一个。应在目标输入长度和并发下观察峰值，并记录是否包含加载、预填充和生成阶段；一次空载成功不能代替这些检查，释放空闲缓存也不能释放仍被张量使用的空间。

记住这三点

  • 权重能载入，只说明初始阶段的一部分空间够用。
  • 普通 KV 缓存随批量、保留长度和实际 KV 维度增长。
  • 训练状态、临时空间与分配器口径，使实际峰值需要单独测量。

别把例子用过头

  • 预算参数完全自拟，程序不创建模型、不分配 GPU，也不报告真实显存测量。
  • 公式只覆盖写明的常规缓存布局，不适用于所有架构、共享、窗口或卸载策略。
  • MiB 与 MB 单位不同；两项小计不是总显存需求。

自己试一小步

怎样解释同一模型在不同请求下的显存变化？

  1. 固定其余假设，分别只改变批量与缓存长度，记录两项小计。
  2. 把缓存元素字节数和权重字节数分开改变，观察哪一项受到影响。
  3. 补一列“公式未包含的空间”，列出实际测量还需记录的阶段和口径。

做到这些，就算完成

  • 没有把查询头数直接当成所有模型的 KV 头数。
  • 相同批量中权重只在本例假设下计一份。
  • 结论没有把估算小计说成硬件可运行保证。

查原始资料

想确认一个细节，先看这份资料对应的部分，不必从头读完所有文档。

找一份资料继续练

先选一项，按下面的范围继续学；不必从头读完整门课程。