大模型显存去哪儿了?一张白板算清训练与推理的容量账本
发布时间:2026/10/18 14:50:13 锦皓数字建站

这一讲我们聊一个所有跑过大模型训练和长文本推理的人都会头疼的问题显存到底去哪儿了。标题起得有点正式实际核心就是一句话——当你盯着nvidia-smi的占用曲线发呆时有没有想过这个数字是由哪些部分组成的。很多人把显存不够等同于模型太大然后开始换小模型、减 batch、砍序列长度调来调去还是 OOM问题在于从头到尾都没把账算清楚。这篇文章不是某类框架的搬运教程而是一套可以拿白板手推的显存与算力账本。读完你能回答三个问题训练一次迭代钱具体花在哪些张量上前向传播一个 token到底做了多少次乘加推理序列拉长到 4K、32K、128K 时KV Cache 什么时候会吃掉比权重还多的显存。最后我们把这些账合并做一次完整的容量规划演练。1. 先看一眼总账7B 模型一次迭代的显存从哪里爆掉1.1 模型状态三件套权重、梯度、优化器谁才是吃显存的大户先统一口径。我们说一个 7B 模型指的是参数总量大约 70 亿个按 BF16 精度存储每个参数 2 字节那么光权重就是 14GB。这是很多人心里的那本账模型 14GB卡有 80GB足够了。但实际一训练就爆显存原因是你只算了第一项。全参数微调或者从头训练时显存里至少要同时住下三类东西模型权重、梯度、优化器状态。前两项很好理解真正吓人的是优化器。以最常用的 AdamW 混合精度训练为例每个参数对应的优化器状态包括一份 fp32 精度的权重主副本、一份 fp32 的一阶动量 m、一份 fp32 的二阶动量 v。每份 4 字节三项加起来就是 12 字节。再加上 BF16 权重副本 2 字节、BF16 梯度副本 2 字节每个参数合计约 16 字节。算一笔具体账7B 模型模型状态 7000M × 16 字节 ≈ 112GB。仅此一项就已经超过单张 80GB 卡的上限还没算中间激活值。所以回答标题的问题显存去哪儿了大头不是权重是 Adam 里那两份动量和一份 fp32 参数副本。这也是为什么很多人第一次用 ZeRO 或者 AdamW 离线统计时会愣一下——权重只占不到八分之一其余全是训练状态。相比之下如果只用 SGD 这类只有一份动量甚至没有动量的优化器优化器状态会大幅缩水模型状态能降到 28GB 左右这就是为什么某些早期模型用 SGD 时单卡显得很“宽松”的底层原因。1.2 中间激活被低估的临时演员训练阶段的第二笔开销是中间激活这部分经常被忽略因为它不像权重那样有个固定的“模型大小”可以查而是严格依赖于 batch size、序列长度和层数。激活的本质是前向传播过程中每个张量都要保留一份副本供反向传播计算梯度时使用。你可以在前向之后立刻删除它们但后向拿什么去求导这就是显存和时间的经典取舍保留得越多后向越省算力但显存越紧张。我以 7B、单条样本、序列长度 2048 为例粗略看一下激活的量级。Attention 里那个 logits 矩阵也就是 Q 和 K 相乘的结果形状是 batch×头数×序列×序列单样本 32 个头就是 32×2048×2048 个元素按 BF16 算约 268MB。这只是最深层的一个矩阵。每个 Transformer 块在反向时还需要保存输入、QKV 投影结果、注意力输出、FFN 中间层输出等一堆张量常见实现里每个块大约要保留十几份形状为 S×H 的中间结果单块粗算 200MB 上下32 层累计就是 6 到 7GB。加上 logits、dropout mask、softmax 统计量等总激活轻松突破 10GB。这里有个非常常见的误解以为激活只看模型维度其实它是“每层都要存一份”。层数越多激活越大序列越长激活越大。想要缓解用得最多的手段是激活检查点也叫梯度检查点前向时把大部分激活丢掉只保留块输入这一层反向时再用保存的输入重新算一遍前向拿回中间值。代价是训练时间上涨 20% 到 30% 左右换来的往往是激活从 10GB 量级降到 3 到 4GB 量级。我自己的习惯是只要显存还够就先不开 checkpoint因为它毕竟是拿时间换空间但一旦 OOM 又不想减 batch这个开关往往是比减小序列长度更优先的选择。1.3 推理阶段的显存构成没有梯度之后钱花在哪推理阶段没有梯度也没有优化器按理说账应该好算很多但实际部署时仍然有显存压力因为钱从“模型状态”转移到了“KV Cache”和“临时计算缓冲”上。先看静态部分7B BF16 权重 14GB 固定不变这部分是刚性的。动态部分有两类一类是每一条请求独享的 KV Cache另一类是算 Attention 时临时创建但会在算子结束立即释放的中间张量比如 QK^T 的 logits 矩阵、softmax 结果、加权求和后的输出。后者在短序列下不值一提在长序列下会突然变得很扎眼因为它的规模随序列长度的平方增长。举个例子序列长度 2048、32 个头时单请求的 logits 矩阵 268MB这个数对 80GB 卡来说不算什么但如果你把序列拉长到 32K同样是 32 个头logits 变成 32×32768×32768×2 字节 ≈ 64GB。这么巨大的临时矩阵显然不能直接落显存所以长上下文推理系统里普遍采用分段计算、分块 softmax 这类手段把超大中间张量拆成小块来处理。这也是为什么标题里要把 KV Cache 单独拎出来讲——它是推理阶段唯一会随并发数和序列长度线性甚至超线性扩张、而且在请求生命周期内一直占住不释放的部分容量规划的核心就是在它身上做文章。2. 白板手算从一次矩阵乘法到 6 倍参数量的 FLOPs 账2.1 从乘法次数开始为什么一个线性层是 2mnp讲显存不能只讲空间还得讲时间。这张卡买回来算力利用率到底怎么样是需要能亲手算出来的。FLOPs 这个概念全称是浮点运算次数理解它只需要一个问题一次矩阵乘法到底做了多少次计算。我们看一个最简单的例子把形状为 m×n 的矩阵 A 和形状为 n×p 的矩阵 B 相乘得到 m×p 的矩阵 C。C 里每一个元素都是 A 的一整行和 B 的一整列做点积一次点积需要 n 次乘法和 n 次加法。C 一共有 m×p 个元素所以总计算量是 (m×p)×n 次乘法加上同样多的加法合计 2mnp 次浮点运算。这个 2 倍因子是整套账本的地基。线性层本质上就是一次矩阵乘法输入形状是序列长度乘隐藏维度 [S,H]权重形状是 [H,H]那么一次前向线性层的 FLOPs 就是 2×S×H×H。这个数字在后面的手算里会反复出现所以我建议你先记住这个结论而不是背公式矩阵乘法的代价大约是 mnp 的两倍。要建立直觉的话可以这么想一个参数量 7B 的模型前向跑一个 token差不多要算 14GFLOPs 的算术运算也就是 140 亿次乘法和 140 亿次加法这个数量级几乎等于每个参数都要参与两次运算。2.2 摊开一个 Transformer 层前向的 FLOPs 怎么累加现在我们把一个标准的 Decoder-Only Transformer 层摊开来看。假设隐藏维度是 H序列长度是 S采用 MHA 结构Q 头和 KV 头数量先统一按 H 维度处理。一个层里主要算子包括QKV 投影、Attention 内部运算、输出投影、FFN 两个线性层。QKV 投影是三份独立的 [S,H]×[H,H] 线性变换合计 3×2×S×H×H 6SH²。Attention 内部运算分两段第一段是 Q×K^T形状是 [S,H]×[H,S]FLOPs 为 2S²H第二段是权重矩阵乘以 V同样 2S²H合计 4S²H。输出投影又是一次 [S,H]×[H,H] 线性层2SH²。FFN 如果有 4 倍隐藏维度的扩张那就是两次线性层分别是 [S,H]×[H,4H] 和 [S,4H]×[H,H]各 8SH²合计 16SH²。把这些全部加起来单层前向 FLOPs ≈ 24SH² 4S²H。如果忽略后面那个与序列平方相关的项只看第一项参数量 N 和层数 L、隐藏维度 H 的关系大约是 N ≈ 12LH²那么 24SH² 就等于 2SN。所以整个模型前向跑完 S 个 token总 FLOPs ≈ 2SN也就是“每个 token、每个参数、两次浮点运算”。这个结论就是行业里常说的“前向 ×2、训练 ×6”的来历训练包含前向一次和反向两次所以每个 token 每个参数大约是 6 次浮点运算总 FLOPs 6×N×T其中 T 是训练总 token 数。2.3 backward 为什么恰好是 forward 的两倍初学者常常觉得反向传播不就是把前向顺序倒过来吗为什么算力是前向的两倍而不是一倍。这里要回到矩阵乘法本身。反向传播对一个线性层 Y XW 要做两件事第一计算 loss 对权重 W 的梯度这需要把输入 X 的转置和上游梯度相乘形状是 [H,S]×[S,H]又是一次完整的矩阵乘法第二计算 loss 对输入 X 的梯度继续往上一层回传这需要把上游梯度和权重转置相乘也是 [S,H]×[H,H]。两件事都需要一次矩阵乘法而前向只需要一次。所以线性层这一层本身反向的算力就是前向的两倍。延展到整个网络每层的前向算一遍后向要为参数梯度和输入梯度各算一遍整体就是前向加上两倍的反向计算总训练 3×前向 6×N×T。这个“6”因子在预训练界是个非常著名的经验常数。你随便翻开一个大模型训练报告算总算力时基本都用这个系数。它不要求你精确到某个具体 kernel 的实现细节但在规划训练时间是够用的。比如训练一个 7B 模型跑 1T token总 FLOPs 6×7×10⁹×10¹² 4.2×10¹⁹如果用一张峰值 990 TFLOPS 的卡理论上需要在 100% 利用率下跑 42000 秒也就是大约 12 小时。但现实里利用率不可能到 100%这就引出了下一个概念。2.4 算完 FLOPs 之后用 MFU 判断你的 GPU 在不在摸鱼MFU模型算力利用率公式特别简单实际算出来的 FLOPs 除以 GPU 峰值算力再除以运行时间。这个指标是我看一个训练系统好不好用的首要指标比盯着nvidia-smi上的显卡利用率准得多。nvidia-smi里那个 99% 只能说明 kernel 没闲着但没闲着不代表在干正事它可能在等数据、等通信、反复读写显存。用上面的 7B 例子验证一下训练每 token 是 42GFLOPs如果一张卡观测到单卡吞吐大约 8000 token/s那实际算力就是 42×10⁹×8000 336 TFLOPS。假设硬件峰值 990 TFLOPSMFU 大约是 34%。这个数字在现代大模型训练里算中等偏上优化得好的系统能做到 40% 到 50%再往上很困难。真正让人意外的是推理场景比如单条请求逐 token 生成时每 token 只有前向的 14GFLOPs算下来 MFU 常常只有 1% 到 2%。不是因为 GPU 不行而是 decode 阶段一个 batch 里只有一个序列权重 14GB 要从显存搬到计算单元这个过程受限于显存带宽而不是算力。带宽决定了一次搬完权重的最短时间也决定了 token 生成的理论上限。这就是为什么长上下文推理系统要想高吞吐必须靠加大 batch 把权重搬运成本摊薄到多个请求头上单纯抱怨“为什么单个请求生成这么慢”是没有意义的问题在账的结构上。3. KV Cache 容量规划长上下文的显存黑洞是怎么算出来的3.1 KV Cache 那个公式2 × 层数 × KV头数 × 头维度 × 序列长 × 精度推理阶段每个 token 进来Attention 都要拿 Q 和所有历史 token 的 K、V 做计算。为了不重复算前面那些 token 的 K 和 V业界习惯把它们缓存下来这就是 KV Cache 名字的由来。很多教程讲到这里就结束了但真正做容量规划时要会算它占多少显存。公式不复杂拆开看就是六个因素相乘每个 token 每层需要为一个 KV 头缓存一份 K 和一份 V所以有个因子 2层数 L每一层 KV 头的数量每个头向量的维度当前序列长度精度字节数。写成直观算式就是KV Cache 字节数 2×层数×KV头数×头维度×序列长度×精度字节数。需要特别提醒的是KV 头数和 Q 头数是两个概念。MHA 结构里每个 Q 头配一个 KV 头KV 头数等于 Q 头数典型 32 个头GQA 结构则可以让多个 Q 头共享一组 KV 头KV 头数可以降到 8 个。这一个区别直接让 KV Cache 缩到四分之一后面对比时你会看到它有多关键。3.2 算几笔具体账GQA 和 MHA 在 4K、32K、128K 下的表现我以 7B 模型、32 层、头维度 128、BF16 精度为例来算。先把“每个 token 每层”的单位成本算出来K 和 V 两份、8 个 KV 头、每个头 128 维、每个维度 2 字节乘起来一层一个 token 是 4096 字节乘以 32 层就是 131072 字节也就是 128KB/token 左右。然后乘以序列长度就比较直观了下面的表格可以直接拿去做预算。序列长度GQA8 个 KV 头单请求 KV 占用MHA32 个 KV 头单请求 KV 占用4K512MB2GB32K4GB16GB128K16GB64GB看到这个表我第一次也愣了一下128K 上下文在 MHA 结构下仅仅是 KV Cache 就是 64GB比模型权重 14GB 高出四倍多。所以在长上下文场景下把“推理显存 ≈ 权重 KV Cache”这句话放在一起看结论会非常颠覆当序列长度足够长KV Cache 才是真正的显存大头。这也说明了为什么现代长上下文模型都在转 GQA不只是为了推理速度更是为了显存容量能撑住更大 batch 和更长序列。如果你用的模型结构是 MHA上下文又习惯拉到几十 K那系统设计上就必须考虑把 KV Cache 放到 CPU 内存做分层卸载或者对 KV 做量化。3.3 容量规划的正确姿势权重、KV、临时缓冲三层预算做推理服务容量规划我会把显存预算拆成三层来估算避免被单一数字带偏。第一层是权重7B BF16 固定 14GB如果做 AWQ 或 GPTQ 4bit 量化可以降到大约 4GB但精度损失要先做评估。第二层是 KV Cache按照上面的公式估算后还要乘以并发数和归一化后的平均序列长度。第三层是临时缓冲包括 Attention 算子内部的 logits、softmax 结果、以及各类 kernel 的 workspace短序列下预留几百 MB 到一两 GB 足够长序列下建议预留总量的一成左右做余量。最终判断一张卡能扛多少个并发请求核心公式是单卡可用显存 权重 并发数×单请求 KV Cache 临时缓冲。假设一张 80GB 卡跑 32K 上下文的 GQA 模型权重 14GB单请求 KV 4GB预留临时缓冲 2GB那么 KV 这一项可用空间是 64GB并发上限 16 路左右。如果是 MHA 模型单请求 KV 变成 16GB并发上限立刻掉到 4 路。同样是 80GB 卡同样模型规模只因为 KV 头数不同吞吐就差 4 倍这就是容量规划要算到底的价值所在。3.4 省 KV 的现实手段GQA、量化、窗口注意力、PagedAttention容量规划不只是算账还要知道怎么调整账本。排在最前面的手段在模型结构层面优先选 GQA 或 MQA 的模型。MQA 是所有头共享一组 KVKV 降到 MHA 的 1/32但效果可能打点折扣GQA 是折中8 个 KV 头换取 1/4 的 KV 占用实际效果在多数任务里和 MHA 差距很小。对齐到工程层面最常用的是 KV 量化。把 KV Cache 从 BF16 压到 INT8KV 占用直接减半压到 INT4 则是减到四分之一代价是精度风险需要配合按通道量化和离线校准来缓解。少有人提的是窗口注意力也就是只保留最近 N 个 token 的 KV远处 token 不参与注意力计算。它跟 GQA 不一样不是省乘法的 KV 数量而是从时间维度上限制缓存长度适合对远程依赖不那么敏感的场景。还有一类是稀疏注意力但实现复杂度偏高不是首选。最后提一个容易踩坑的细节KV Cache 是随请求动态增长的如果预分配整块显存按最大长度算内存碎片会很严重容量利用率直线下降。实际系统普遍采用按页分配的方式类似操作系统内存分页需要新增 token 时再申请新页而不是一次性锁死一大块连续空间。分页字段之间还可能因为页大小配置不当造成浪费这个我后面会专门讲一个具体的踩坑例子。4. 把账本合起来一次实际部署案例的容量规划演练4.1 场景设定7B 全参训练 / LoRA 微调 / 32K 上下文推理前面两章分别讲了训练显存、FLOPs 和 KV Cache现在把它们合到同一个项目里做一次完整的容量规划演练。假设我们拿到一个 7B 开源模型结构是 32 层、隐藏维 4096、32 个 Q 头、8 个 KV 头GQA、BF16。要做的任务是三件事第一用全参数微调在一批领域数据上做继续训练第二用 LoRA 做一套更轻量的微调方案第三部署一套支持 32K 上下文的高并发推理服务。看起来是三个独立任务但它们共用同一本底账思路完全一致先算权重再算训练状态或 KV最后看还剩多少显存给 batch 和并发。为什么把三个任务放一起讲因为太常见了。你白天调训练脚本被 OOM 折磨晚上部署推理时又发现并发上不去问题表面不相关根子里都是同一个原因没有按前面那套思路估算显存结构而是靠试错调整 batch。试错不是不行但每次 OOM 都要往前回溯好几个变量浪费的时间太多了。4.2 训练方案为什么全参必须 ZeROLoRA 为什么只省一半先算全参微调。模型权重 14GB梯度 14GBAdam 状态 84GB合计 112GB。单卡 80GB 肯定装不下所以必须上 ZeRO 或张量并行。以 ZeRO-3 为例优化器状态、梯度、权重都可以按 rank 分片单卡承担的模型状态大约等于 112GB 除以 GPU 数再加上激活。比如 4 卡每卡模型状态约 28GB留出激活 8GB40GB 左右80GB 卡跑起来很稳。这里不要只看“能不能跑”还要看通信开销ZeRO-3 在每次前向后向都要做权重收集和梯度规约通信量大所以卡间互联带宽不够时会明显拖慢训练速度这也是为什么 ZeRO-3 在高性能集群里并不是无脑最优。再看 LoRA。很多人误以为 LoRA 省显存是因为冻结了原模型重量其实权重 14GB 一分不少照样要放显存LoRA 真正省的是梯度和优化器状态。因为 LoRA 只训练注入的低秩矩阵梯度只有几十上百 MBAdam 状态只针对那几十 M 参数模型状态从 112GB 直接降到约 15GB。加上激活 8GB 左右单张 80GB 卡完全够用。所以 LoRA 不是“省钱版模型”而是“省训练状态”的微调方式这个认知差异决定了你能否精准定位 OOM 瓶颈如果跑 LoRA 还 OOM那问题一定不在优化器状态而在权重复制、激活或者缓存没清干净。4.3 推理方案KV Cache 和权重谁决定你的 batch size推理任务按 32K 上下文、GQA 模型来算。权重 14GB 固定单请求 KV Cache 按前面的表是 4GB临时缓冲预留 2GB。一张 80GB 卡在 32K 上下文下最多能支撑大约 16 路并发请求。而且这 16 路是理论极限实际要留显存给 CUDA context 和框架运行时一般再打个八折也就是 12 到 13 路。如果想上更大的并发三个方向把权重量化到 4bit权重降到 4GB多出 10GB把 KV Cache 从 BF16 压到 INT8单请求 KV 从 4GB 变 2GB或者缩短平均序列长度比如绝大多数请求只用到 8K 上下文就按 8K 预留而不是按最大预算。这套账算完之后你会发现一个反直觉的结论在这个例子里KV Cache 决定了批大小权重反而成了次要因素。只有当你把 KV 压缩到 1/4 以下、权重量化降到 4bit权重和 KV 才重新回到差不多的量级。这也是为什么业界做长上下文推理时普遍把稀疏注意力、KV 量化和权重量化当套餐用而不是只押注其中一项因为单一手段省下来的空间很快就会被翻倍的并发需求吃回去。我在实际推进类似部署时习惯先做一个 30 分钟的压力测试固定权重和 KV 精度连续增加并发请求观测显存增长曲线和首 token 延迟验证纸面估算和实测的偏差。如果实测显存明显高于估算优先怀疑两个地方一是显存碎片二是框架为每条请求预分配的 buffer 过大。这两类问题在账本上算不出来只能在真实负载里看到。5. 几个容易看走眼的显存细节和我的处理习惯讲完大账之后再分享几个我在真实项目里踩过、后来总结成固定检查项的坑。第一个是优化器状态的口径问题。很多框架的显存记录把 Adam 里的 fp32 master weight、momentum、variance 分开显示但文档里写的“每个参数 12 字节”是合并口径如果你只把其中一项当成全部预算会差三倍。我现在的做法是统一按参数总量乘 16 字节来粗估 AdamW 训练状态不管框架界面怎么显示。第二个坑藏在 KV Cache 的分页分配里。分页的页大小如果设置得比 model 实际每 token KV 消耗小会频繁申请新页如果设得过大每个页尾部的空闲空间就是浪费。之前遇到一个线上服务显存利用率虚高排查到最后是页大小设为 64 个 token而 KV 平均请求长度只有 300 token 左右尾部浪费接近 10%。从那以后我都会先按模型的单 token KV 字节数计算这个模型在某个页大小时实际能装多少 token再反推显存浪费率。页大小宁愿设置为 4 到 8 个 token 的整数倍也不要盲目设成很大因为 KV Cache 是动态增长的页太大会造成大量只填充了不到一半的空洞。第三个习惯是关于量化精度的选择。KV 量化从 BF16 压到 INT8对显存的效果立竿见影但如果你用的是按静态缩放因子做均匀量化不在线更新 scale当序列长度分布非常不均匀时会出问题。我碰到过一次个别超长请求在同一层同一头里出现了极大值INT8 量化后表现异常最后换成分组量化才稳定下来。所以我现在的建议是做 KV 量化前先拿真实请求分布做一次离线校准不要拿一个简单的“KV 压缩一半”来决定方案精度和容量是同一本账的左右两页。最后回到这一讲的主题词显存全景账本、白板手算、KV Cache 规划。这套东西学完本质上练的是一种脑内建模能力——拿到一个模型先能推算出它每个阶段大概吃什么量级的资源再带着这个判断去看工具文档而不是被工具文档牵着走。我自己做训练和部署这么久最深的体会是几乎所有显存相关的问题都能归结为“没分清楚哪份账该由哪一层来承担”。模型状态归优化器激活归激活检查点KV Cache 归结构设计和量化临时缓冲归算子实现。每一层都按这个原则去优化显存不够的问题往往比想象中更容易解决。
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。