搞深度学习这些年被问得最多的一句话是明明买的 80G 的卡模型还没跑两步就 OOM显存到底去哪了 到了大模型训练和推理阶段这个问题更扎心一个 7B 模型 BF16 权重才 14GB为什么全参训练动不动就要上百 GB为什么推理时上下文一拉长KV Cache 转眼就吃掉几十 GB要回答这些不能靠感觉得会算账。这篇就当第 27 讲把显存全景账本、FLOPs 白板手算、KV Cache 容量规划三件事一次讲透训练和推理的钱分别花在哪几项矩阵乘法到 Transformer 的 FLOPs 怎么手推KV Cache 的公式从哪来、不同模型吃多少、并发怎么规划。适合正在做微调、部署推理服务或者总被 OOM 折磨的同学看完你就能在白板上把显存需求算明白而不是等任务跑挂了才后悔。1. 显存全景账本训练和推理的钱分别花在哪1.1 训练账本权重、梯度、优化器状态与中间激活先说训练。很多人只盯着模型权重看觉得 7B 模型 14GB 权重一张 24G 卡总够了吧实际完全不是这样。显卡在训练时要同时装下四类东西模型权重BF16/FP16 下每参数 2 字节7B 就是 14GB梯度和权重同尺寸也是每参数 2 字节14GB优化器状态AdamW 要维护一阶动量 m 和二阶动量 v各自是 FP32、每参数 4 字节合起来 8 字节/参数56GBFP32 主权重混合精度训练里为了避免 BF16 精度不足导致更新失效框架通常还会额外保留一份 FP32 主权重每参数 4 字节28GB。把这几项加起来228416 字节/参数7B 模型就是 112GB。这就是为什么 80G 单卡根本装不下 7B 全参微调必须用 ZeRO/FSDP 把梯度、优化器状态分片到多卡。这里有个经常被忽略的细节如果不保留 FP32 主权重每参数可以省 4 字节总占用降到 84GB228。但代价是 Adam 更新收敛变差特别是小学习率场景很容易踩精度坑。我建议你把 16 字节/参数当作训练内存估算的标准默认值只有明确知道自己在做什么时才往下降。另一个思路是 LoRA/QLoRA冻结原权重只训练低秩适配器优化器状态只按 rank 维度的少量参数计算成本能降一个数量级这也是显存不够时最常用的出路。除了这几大块还有一笔容易被低估的开销中间激活。反向传播要重算梯度就必须把前向过程里每一层的输入、输出、dropout mask 留住。经验公式是不做激活重计算activation checkpointing时每层大约要 B×S×H×34 字节BF16 场景也就是 batch、序列长度、隐层维度三者相乘再乘一个常数。拿 B4、S4096、H4096、L32 举例每层约 2.3GB32 层就是 73GB。开了激活重计算后这笔钱可以压到 1~2GB 级别代价是前向要多算一遍FLOPs 涨 30% 左右。等后面讲 FLOPs 时你会看到这笔交易到底划不划算。注意上面的 34 是工程上的经验常数不同框架、不同算子融合程度会有差异。你要做的是用它做数量级判断真正落地时以框架 profiler 的数据为准。1.2 推理账本权重、KV Cache 与临时张量推理和训练完全是另一套账。推理不需要梯度、不需要优化器状态主要开销是三项权重照旧是每参数 2 字节BF16或量化后更低KV Cache自回归生成时必须记住历史 token 的 Key 和 Value后文会专章推导先记住它是随上下文长度线性增长的活钱临时张量attention softmax 的中间结果、gating 输出等prefill 阶段比较占decode 阶段通常很小。很多同学问为什么训练把激活算一次反而不够用推理反而要单独为 KV Cache 做容量规划因为训练是一锤子买卖激活是短暂的推理的 KV Cache 却要持续占着直到该序列被淘汰。预填充prefill阶段模型是算力瓶颈1 个 token 的前向几乎要把所有权重读一遍解码decode阶段则变成典型的内存带宽瓶颈——每生成一个 token 都要重新把全部权重从显存搬到计算单元所以单用户生成速度基本等于权重字节数 / 显存带宽。这也是同样的模型8 卡训起来飞快、单卡推理却很慢的根本原因。推理阶段还有一笔很多人不算的账CUDA 上下文和显存碎片。空卡启动一个进程通常要吃掉 0.5~1GB长时间反复申请释放 KV block 后碎片会越来越严重。所以做容量规划时我给每张卡留 2GB 左右的系统冗余不要算得刚刚好。1.3 用 7B 模型算一笔完整账单把上面两张账合起来拿一个 7B 稠密模型做例子你会看到训练和推理的显存需求完全是两码事项目训练BF16 Adam推理BF16 单路权重14GB14GB梯度14GB无Adam 动量 m/v56GB无FP32 主权重28GB无KV Cache4K 上下文GQA 8训练无此项约 0.5GB中间激活B4, S4K数十 GB 量级约 1~2GB临时张量与系统开销少量约 1~2GB训练这边合计轻松突破 100GB推理单路才 17GB 左右。这个差距解释了所有工程实践里的怪现象训练怕的是优化器状态和激活推理怕的是权重读取带宽和 KV Cache 扩容。对症之后优化方向完全不同训练优先做分布式状态分片和激活重计算推理优先做权重量化、GQA 改造和 KV Cache 管理。很多人拿着训练的思路去优化推理或者反过来结果自然是白忙一场。2. FLOPs 白板手算从矩阵乘法到 Transformer 全套2.1 矩阵乘法的 FLOPs 到底怎么数先立规矩。一次矩阵乘法 C[M,N] A[M,K] × B[K,N]要做 M×N×K 次乘加每次乘加算 1 次乘法和 1 次加法所以业界统一口径是 2×M×N×K FLOPs。这个2是新手最容易漏的。一个线性层 y xW输入 x 是 [B,S,H]权重 W 是 [H,O]一次前向就是 2BSHO FLOPs。按 token 均摊就是每个 token 2HO。记住这个基本单位就够了后面所有 Transformer 的 FLOPs 都是把它套进每一层逐项加出来的。我在白板上推的时候习惯一行一行写而不是直接背公式因为只有推导过一遍你才知道哪个常数从哪来模型结构一改也不会算错。2.2 Transformer 单层前向的推导过程假设模型隐层维度 H标准多头注意力每个 head 维度 d_k H / n_heads。对于单个 token单层前向的开销如下QKV 投影三个矩阵乘法每个是 H×H共 3×2H² 6H²注意力分数 QKᵀ带因果掩码时平均每个 query 要跟约 S/2 个 key 做点积每个 head 开销约 S×d_k所有 head 合计约 H×S注意力加权求和 AV与 QKᵀ 量级相当也是约 H×S输出投影一个 H×H2H²MLP标准两层、中间 4H第一层 H→4H 是 8H²第二层 4H→H 也是 8H²合计 16H²若是 SwiGLU 这类三线性层结构则要按 24H² 计。逐项相加标准 MLP 结构下单层前向约 24H² 2HS 若干可忽略项。细心的人会发现24H² 恰好等于 2 × 12H²而 12H² 正是这一层的参数量QKV 3H² 输出投影 H² MLP 两层 8H²。这说明一个很好用的结论不看注意力那点序列相关项的话每层前向 FLOPs ≈ 2 × 该层参数量整个模型前向 1 个 token ≈ 2 × 全部参数量。所以 7B 模型生成 1 个 token纯算力上约 14 GFLOPs8K 上下文的 prefill约 2×7B×8192 ≈ 1.15×10¹⁴ FLOPs也就是 115 TFLOPs。这个量级配合高端加速卡的纸面 BF16 算力prefill 应该是百毫秒级——但实际往往还要加上注意力、通信、内存搬运等开销。等看完 decode 阶段你就明白算得快不代表生成快瓶颈经常根本不在 FLOPs 上。提示QKᵀ 和 AV 两项里的 S 是序列长度当上下文拉长到几万甚至百万 token 时这部分的增长会让FLOPs ≈ 2×参数量的近似失效。长上下文场景记得把这两项单算否则你会低估总算力需求。2.3 训练为什么是 6 倍前向反向传播的 FLOPs 拆解反向传播凭什么比前向贵一倍以一个线性层为例。前向 y xW 只有一次矩阵乘。反向要算两个东西对输入的梯度 dx dy·Wᵀ和对权重的梯度 dW xᵀ·dy。这两次矩阵乘的规模都与前向相同所以反向的 FLOPs 是前向的 2 倍。逐层累加训练一个 token 的总 FLOPs 就是前向 2N 反向 4N 6N也就是常说的训练一个 token 约等于 6 倍参数量 FLOPs。这是大模型训练一切算力规划的地基。有了 6N你可以很快算训练总计算量6 × 7B × 500 亿 token ≈ 2.1×10²¹ FLOPs。假设你用 8 张 A100BF16 纸面约 312 TFLOPS实际训练 MFU 按 40% 估总有效算力约 8×3.12×10¹⁴×0.4 ≈ 10¹⁵ FLOPs/s那么训练时长约 2.1×10²¹ / 10¹⁵ ≈ 2.1×10⁶ 秒约 24 天。这样你在规划数据量和卡数时直接拿 6ND 一除就有数了不用等集群排队跑一轮才知道行不行。再补充一个工程常识开了激活重计算后实际 FLOPs 会从 6N 涨到约 7.2~8N而使用张量并行、流水并行时通信本身不算 FLOPs 但会拖慢有效吞吐。所以真实世界里 MFU 能到 50% 就算很健康了白板手算的理论下限和实际耗时要留出 2~3 倍余量。阶段每个 token 每层的 FLOPs说明前向约 2×参数量线性层为主注意力项随 S 增长反向-输入梯度约 2×参数量dx dy·Wᵀ反向-权重梯度约 2×参数量dW xᵀ·dy训练合计6×参数量前向 反向3. KV Cache 容量规划公式推导与实例测算3.1 KV Cache 大小公式与 GQA 改造KV Cache 的来历要从注意力机制本身讲。自回归生成到第 t 步时第 l 层的 query q_t 要和过去所有 token 的 key 做点积再对这些 key 对应的 value 做加权求和。如果不想每步都重算历史 token 的 K、V就必须把它们存在显存里这就是 KV Cache。它的大小由四个维度决定层数 L、batch 大小 B、已生成长度 S、以及每层 KV 的隐层宽度 H_kv再乘 2K 和 V 各一份、乘上每元素字节数。写成白板公式就是KV Cache 字节数 2 × L × B × S × H_kv × 字节/元素其中 H_kv 由注意力头结构决定MHA多头注意力下 H_kv HGQA分组查询注意力下 H_kv H × n_kv / n_q也就是只有查询头分组共享少数几个 KV 头。之所以说 GQA 是 KV Cache 的主刀医生是因为它直接砍掉的是乘数项里最贵的那一项。MHA 到 MQA 是极端方案所有查询头共享一组 KVKV 占用缩到 1/n_qGQA 是中间态用 2~8 个 KV 头换来质量损失最小的压缩。举两个落地例子。7B 级模型的典型参数L32、H4096、n_q32、d_k128。MHA 下每个 token 每层要存 2×4096×2 字节 16KB整模型每 token 就是 512KBGQA 用 8 个 KV 头后降到 128KB。70B 级模型则更夸张L80、H8192、n_q64MHA 下每 token 约 2.6MBGQA-8 后仍有 320KB。你会发现 KV Cache 的增长几乎和模型规模、上下文长度是线性叠加的关系这也是为什么长上下文系统几乎清一色上 GQA KV 量化。3.2 常见模型的容量测算对比用同一个公式把常见配置的 KV 开销列出来单序列、BF16、不开启 KV 量化模型规模层数/隐层KV 头结构每 token 开销4K 上下文32K 上下文128K 上下文7B 级32 / 4096MHA512KB2GB16GB64GB7B 级32 / 4096GQA-8128KB0.5GB4GB16GB70B 级80 / 8192MHA2.6MB10.5GB86GB344GB70B 级80 / 8192GQA-8320KB1.3GB10.7GB43GB看到差距了吗70B 模型如果还是 MHA跑 128K 上下文时单条对话的 KV Cache 就是 344GB比整机显存都大这显然是不可服务的。而同样 70B 用 GQA-8128K 上下文降到 43GB——仍然很重但至少可以靠多卡张量并行硬扛。如果在 KV Cache 上再叠加 4-bit 量化还能再砍一半代价是精度略有下降。这就是为什么长上下文不是免费的每多一万 token 上下文70B 级模型就要多准备 3.2GB 显存GQA-8你得想清楚自己的业务到底需不需要这么长的上下文。不少团队在 32K 和 128K 之间反复横跳最后发现真正高价值的用户请求平均都不到 8K为此付出十几倍的 KV Cache 成本并不划算。3.3 容量规划实战并发、上下文与量化组合容量规划的本质是一道算术题。以单卡 80GB 部署 7B BF16 推理为例权重 14GB系统冗余留 2GB预留给 KV Cache 的空间约 64GBGQA-8、4K 上下文时每条并发序列约 0.5GB理论上能开 128 路并发同样配置换成 MHA每路变 2GB并发直接掉到 32 路若把上下文拉到 32KGQA-8 下每路 4GB并发只剩 16 路。这个账一算就明白并发数、上下文长度、KV 量化精度三者是跷跷板。部署时我习惯先定业务 SLA 里的最大上下文再反推可接受并发最后才决定要不要上 KV 量化。另外还要预留好 prefill 阶段的临时激活——特别是 batch 大、上下文长的时候prefill 的中间张量可能瞬间吃掉几十 GB虽然只持续几十毫秒但在峰值统计里会造成严重的显存尖峰。建议用框架自带的 profiler 或显存监控把 prefill 峰值和 decode 稳态分开记录不要只看平均占用否则你规划的 128 路并发会被 prefill 的瞬时峰值直接击穿。工程上还有一批 KV Cache 优化技巧可以叠加滑动窗口只保留最近 N 个 token 的 KVtoken 淘汰/压缩策略KV Cache 换出到 CPU/内存做 offload以及高性能推理框架普遍采用的页式 KV 分配消息队列调度、显存碎片管理这些思路都被融合进去了。它们都是在上面公式基础上的进一步抠内存但再花哨的操作也改变不了公式里最基础的几个乘数——规划时先把公式算对再做优化不然只会越优化越混乱。4. 实操中的显存排查与常见问题速查4.1 四步定位显存去哪了看到 OOM 不要慌按顺序排查比瞎调参管用得多。第一步分清是分配还是保留。nvidia-smi显示的是驱动层面的显存占用包含 CUDA context、缓存块等而框架自带的torch.cuda.memory_summary()才显示进程内已分配和未分配的缓存。很多所谓显存泄漏其实是框架缓存块没归还不代表真的泄漏重启进程或调低缓存上限就解决了。第二步区分峰值出现在哪个阶段是前向 rollout、反向梯度回传、优化器更新还是推理的 prefill/decode。可以在关键节点用torch.cuda.reset_peak_memory_stats()和torch.cuda.max_memory_allocated()打点一次训练任务打完点峰值位置一目了然。第三步逐项开关变量做对照测试关闭激活重计算看激活占多少、把 batch 减半看线性比例、只跑一次反向看梯度那部分是否异常。第四步查碎片。如果分配的峰值不高但总是 OOM大概率是显存碎片化严重尤其频繁申请释放 KV block 的推理服务常见此问题页式 KV 分配类方案就是冲这个来的。实操心得遇到训练 OOM第一反应不要查激活公式而是先看打点日志里峰值出现在哪个函数。我见过太多人凭感觉把 batch 缩了一半结果问题出在优化器状态的 56GB 上怎么缩 batch 都没用最后加一张卡才解决。4.2 常见显存问题速查表现象大概率原因快速解法前向过程中 OOMS 和 B 越大越严重中间激活爆炸开激活重计算、减小 batch、用梯度累积反向或优化器更新时 OOM梯度/优化器状态超限ZeRO/FSDP 分片、换 LoRA、去掉 FP32 主权重推理上下文越长越 OOMKV Cache 超限上 GQA 模型、KV 量化、降上下文上限、限并发总占用不高但反复 OOM显存碎片重启进程、用页式 KV 分配、提前统一申请换卡后显存占用变化异常环境变量与物理卡不对应检查容器内的 CUDA_VISIBLE_DEVICES 与nvidia-smi对应关系这张表不是让你背的而是排查时的思维索引。看到 OOM 先对号入座再决定从哪个方向动手能省下一整个下午的盲目调试。我自己的习惯是所有显存问题先用一句话描述峰值出现在哪个阶段 哪个配置项线性放大能写清楚这句话问题基本解决一半。4.3 踩坑心得与显存规划建议最后分享几条我自己踩出来的规律。第一永远按峰值算账、按冗余部署。白板估算出来 17GB 的推理服务我至少按 20GB 规划训练估算 112GB我不会在 2×80GB 上卡线跑而是留出通信和临时 buffer 的空间。第二训练内存用16 字节/参数 激活项打底推理用2 字节/参数 KV 项打底这两句口诀覆盖了我 90% 的规划场景。第三公式算完一定要用 profiler 回验一次模型结构、框架实现细节会让常数有 10%~30% 的偏移只有实测数据才能校准你的手算模型。另外一个很容易踩的坑是序列打包sequence packing。为了提升训练效率把不同长度的样本拼到同一个 batch 里激活的 S 会按最长长度算显存可能瞬间翻倍。如果你发现训练吞吐没升多少、显存却先爆了多半就是打包策略没有按实际长度分桶。推理侧同理批量请求的 KV Cache 是按每条序列实际长度计费的长短请求混在一起时要防止长上下文请求把整卡 quota 吃掉最好按上下文长度分层限流。我在实际项目里还有个习惯任何新模型上线前先跑一个最小化的热身请求把 prefill 峰值、decode 稳态、KV 增长速率三个数据打出来存档作为后续所有容量规划的基准。白板手算解决方向对不对的问题实测数据解决数量准不准的问题两者结合显存管理就不再是玄学。最后再分享一个小技巧别总盯着大模型怎么省显存先问自己怎么花得明白。训练时把优化器状态、梯度、激活三大头分别打点统计推理时把权重、KV、系统开销分开记账每一笔都对应到公式的某个乘数上。这样你会突然发现真正需要优化的不是某个玄学参数而是最便宜最好改的那一两个乘数——对大多数人来说先把 batch 和上下文长度管好显存问题已经解决了八成。
