先上结论这份讲义的真正价值不是教你怎么背出“37倍、16倍”这些系数而是让你在白板上从零推演一遍把模型从参数、激活、梯度到KV Cache的每一项显存开销全部算清楚。它能解决的实际问题是当你的训练任务再次报错“CUDA out of memory”或者推理服务的并发上不去时你不再靠玄学调参而是能在60秒内定位瓶颈是权重、优化器状态、中间激活还是KV Cache并给出一个有依据的量化判断。这篇内容适合三类人一是准备部署大模型、需要做推理资源规划的工程师二是想理解训练显存放不下、需要做模型并行或梯度检查点决策的算法同学三是面试前想搞清楚“为什么7B模型训练至少要56GB显存”这类问题的候选人。我会按显存全景账本、FLOPs白板手算、KV Cache容量规划三个递进层次展开每一步都给出可复算的公式和真实案例。1. 显存全景账本先分清三笔“花销”和三笔“储备”显存管理的第一个误区就是只知道“模型太大放不下”却说不清显存到底被谁吃掉了。要建立全景账本先记住一个原则显存里的一切都可以归结为两类——静态的、加载一次就常驻的和动态的、每个step或每来一个token都要重新计算的。这两类的应对策略完全不同。1.1 三笔花销参数、激活、KV Cache我们用一张“你是谁、占多少、何时释放”的表格把三笔主要开销钉死开销项存储内容量级特征生命周期模型权重每一层矩阵的W、b以及embedding表固定与上下文长度无关加载后常驻优化器状态Adam的动量m、方差v训练时固定通常2~4倍于参数训练期间每step更新中间激活前向传播每层的输出、注意力矩阵、临时结果与batch size × sequence length成正比每个step临时分配反向时用于梯度计算KV Cache推理时缓存的Key/Value序列与并发数和上下文长度成正比随请求生命周期存在很多人会把“激活”和“KV Cache”混在一起这是不对的。训练阶段激活是为主——因为反向传播要拿激活值去求梯度推理阶段KV Cache是主——因为decode阶段每个token都要去查询前面所有token的Key/Value不缓存就每天重复算几万次注意力。这两者不会同时成为主要瓶颈因为你不会一边训练一边推理。1.2 三笔储备权重、梯度、优化器状态训练态显存账本的起点是“Adam优化器 混合精度”这个最标准的组合。以参数量为P的模型为例模型权重本身如果FP32占4P字节如果用FP16/BF16混合精度模型主副本仍是FP32另有一份FP16的实时权重用于前向和反向所以权重部分大约是4P 2P 6P字节。梯度混合精度训练下梯度用FP16存储但Adam更新时需要FP32精度所以实践中梯度也常保留FP32版本占4P字节。优化器状态Adam除了参数本身还要保存动量m和一阶矩估计v两者都是FP32共8P字节。把这四项加起来6P 4P 8P 18P字节。这就是“7B模型训练至少需要126GB显存”这个说法的来源。但请注意这只是训练态的基础盘还没有算激活。激活在最坏情况下会把总需求推高到60P字节以上所以真正决定训练能否跑起来的往往是激活管理而不是参数本身。“6P 4P 8P”这个公式我建议你亲手推一遍模型参数2P半精度权重加上梯度2P、优化器8P再考虑主权重4P。推导过程比背结论有价值——一旦换了优化器比如用SGD优化器状态只有动量一项你能立刻调整账本而不是死记系数。1.3 训练态与推理态的“分水岭”在于显存峰值继续深入训练态。上面说的18P字节是“稳态”的最低要求。真正的显存峰值出现在某个特殊时刻——当batch里的每一个样本在每一层都留下完整激活图的时候。举个例子一个7B模型batch size 1seq len 2048hidden 4096。仅注意力计算中的QK^T矩阵和softmax结果每层就要存8 × 4096 × 2048 × 2048字节按BF16算约67MB几十层累计下来就是几个GB。这还只是注意力部分不包含FFN的中间激活。峰值显存和均值显存是两个概念很多OOM发生在某个瞬时峰值上而你监控到的平均值可能完全正常。推理态的显存账本就清爽得多。推理不做反向传播没有梯度也没有优化器状态权重可以全部用FP16/BF16加载所以基础需求只有2P字节。再加上KV Cache和临时激活才是推理进程的真实占用。因此同样的7B模型推理态的显存基线是14GBBF16而不是126GB。这也是为什么在直觉上训练要8卡A100推理只要1~2张卡就可以开始玩的原因。1.4 显存规划的先后顺序先静态后动态结合多年跑模型的经验我建议你在做任何容量规划时都按这个顺序先把静态开销列出来权重 优化器状态 保留buffer。再估算动态开销激活或KV Cache。对比你的卡容量看瓶颈落在哪一侧。再决定用哪种策略静态超了考虑模型并行/量化动态超了考虑梯度检查点/更小的batch/限制并发。这套顺序看起来简单但很多人一开始就扎进“KV Cache怎么调”的细节里结果权重放不下也浑然不觉。先静态后动态能帮你快速定位问题层级。2. FLOPs白板手算从矩阵乘法开始推不背公式显存是“盘子”FLOPs是“锅里的菜”。想验证显存账本是否合理最好的方法是用计算量来交叉校验如果你的模型理论计算量是某张卡理论算力的40%但实测利用率只有20%那说明显存或内存搬运成了瓶颈。这种交叉校验能力来自你对FLOPs的白板手算能力。2.1 一次矩阵乘法到底要算多少次乘加一条最基础也最容易被忽略的规则矩阵AM×N和矩阵BN×P相乘结果矩阵是M×P每个输出元素要做N次乘法和N次加法。通常我们把一次乘加MAC记为2 FLOPs所以总计算量是 2 × M × N × P FLOPs。举例A是(4096, 4096)B是(4096, 4096)那么 2 × 4096 × 4096 × 4096 ≈ 137.4 GFLOPs。一张A100的理论FP16算力约312 TFLOPs带稀疏则翻倍理论上这个矩阵乘法只需0.44毫秒。但实际如果显存带宽受限可能耗时翻好几倍。这就是手算FLOPs的价值——你可以判断自己到底是被算力卡住还是被显存带宽卡住。2.2 逐层拆解Transformer的计算量标准的解码Transformer层包含四类主要计算QKV投影、注意力打分与加权、输出投影、FFN的两个全连接层。我们假设hidden size hseq len sbatch bFFN中间维度 4h。以第l层为例QKV投影输入xb×s×h分别和三个矩阵h×h相乘总FLOPs 2 × b × s × h × 3 × h 6bsh^2。注意力打分Qb×h×s×头数维度与K转置相乘一次是2 × b × s × s × 内部维度内部维度为head_dim总head数为h/head_dim所以合起来是 2bsh × s 2bs^2 h。注意力加权attn_scores与V相乘公式与上一步完全对称也是2bs^2 h。输出投影attention输出和W_oh×h相乘FLOPs 2bsh^2。FFN两层先升维FLOPs 2 × b × s × h × 4h 8bsh^2再降维同样8bsh^2。把以上相加得到单层Transformer的计算量约为 24bsh^2 4bs^2h。其中注意力部分的 4bs^2h 只在s较大时有感比如s4096时注意力占比明显上升。这也是现在强调稀疏注意力、线性注意力的原因之一。2.3 为什么手算FLOPs有两个常犯的坑第一个坑是“把FLOPs和Param算混淆”。参数计算只看矩阵维度运算量计算还要乘上batch和seq len。有些文章说“GPT-3的FLOPs是3.14×10^23”你需要理解那是把所有token、所有层、所有batch都乘进去的总量不是单层单样本的量。第二个坑是“漏掉反向传播”。训练的前向FLOPs只是总FLOPs的一部分。标准经验是反向传播需要的计算量约为前向的2倍原因是一次反向需要先对损失求梯度再对输入求梯度而实际分解起来约等于权重梯度的计算量加输入梯度的计算量。所以训练一层的总FLOPs ≈ 3倍前向FLOPs。如果你用前向FLOPs去估算训练时间结果通常对不上号。手算FLOPs的经验法则是先用一层算单层量再乘以层数得到一次前向量训练再乘3。这个流程能让你在10分钟内对比不同模型规模、不同序列长度下的计算压力变化曲线比任何论文里别人给的估算都更贴合你的实际配置。3. KV Cache容量规划一个公式说清推理显存的“大头”如果说训练显存的主角是激活那推理显存的主角就是KV Cache。在自回归解码阶段每生成一个token模型都要拿当前的Query去查所有历史token的Key和Value。如果不缓存时间复杂度就是平方级别增长缓存之后每次生成只需要算当前token的注意力复杂度降到线性。KV Cache就是用显存换算力的经典做法。3.1 KV Cache的一个token占多少显存推出KV Cache容量公式前先明确几个符号L为层数h为hidden sizek为GQA/MLA的KV头数因子如果每层的KV头数与Q头数相同则k 1如果使用分组查询注意力GQAK和V的头数可能是Q头数的1/2、1/4等k会小于1precision为每个数占的字节数s是上下文长度b是并发请求数。每个token在每个层、每个注意力头中都要存储一个Key向量和一个Value向量。假设KV头数总共为n_kv_head每个头的维度为d_head那么每个token在单层要存的KV字节数为 2 × n_kv_head × d_head × precision。因为n_kv_head × d_head hidden所以单层每个token的KV存储为 2 × hidden × precision字节。把这个结果扩展到全部层和全部token每个token的KV Cache大小 L × 2 × hidden × precision 字节。对于常见配置比如L32hidden4096BF16precision2每个token占 32 × 2 × 4096 × 2 512KB。看着很小但上下文长度s8192时单条请求的KV Cache就是512KB × 8192 4GB如果并发8路就是32GB。这个结果往往会吓人一跳因为很多工程师以为推理显存主要是权重占的实际长上下文高并发下KV Cache才是绝对的大头。这里要特别提醒如果你的模型使用了GQAKV头数只有Q头数的1/2或1/4那么公式中的“hidden”要替换成“KV投影维度的总和”。换个说法GQA并不是把缓存变成“hidden的一半”那么简单要直接看模型的KV投影输出维度。这也是很多实现中“KV Cache计算和理论对不上”的原因。3.2 容量规划实例7B模型的推理到底吃多少显存我拿一个真实场景来算模型是L32hidden4096的7.1BBF16权重context4096并发32路。权重BF167.1 × 2 14.2GB。KV Cache每token32 × 2 × 4096 × 2 512KB。每请求context 4096的KV Cache512KB × 4096 2GB。并发32路的总KV Cache2GB × 32 64GB。推理总显存 ≈ 14.2 64 激活buffer ≈ 80GB上下。这就是为什么部署7B模型做长上下文高并发服务动辄要上A100 80GB甚至多卡的原因——多数人以为14.2GB权重就够了实际算完KV Cache后空间就翻了好几倍。建议训练一个自己的“随手估”习惯权重约2PKV Cache约L × 2 × h × precision × tokens。两个公式加起来任何模型推理要多少显存都能在10秒内给出一个数量级正确的答案。如果和实际观测差太多优先检查KV Cache的量化位数和GQA配置。3.3 KV Cache的量化、吸收和复用KV Cache容量大的解决方案业内已经有几条成熟路线但我建议无论选哪条都要先理解公式里的哪个变量被动了。KV Cache量化把precision从2字节降到1字节甚至0.5字节INT8/FP8或4位量化KV Cache直接减半或减到1/4。代价是精度损失需要校准确认。GQA/MQA从模型结构上减少n_kv_head每token的KV Cache变小。这是训练阶段就要定好的架构选择推理阶段无法改。PagedAttention解决的是KV Cache碎片化和内存利用率问题不是减少总量而是让显存中零散的小块也能被利用起来。它的核心是把KV Cache按固定大小的块分配就像操作系统的分页。这条技术对推理引擎优化非常关键值得单独研究。跨请求缓存比如System Prompt共享时可以减少重复计算和缓存占用。这种优化在RAG场景下收益很大但工程复杂度不低。3.4 推理态完整账本权重 KV Cache 激活这里的“激活”其实是一个小角色。在decode阶段由于每次只生成一个token激活值只需按单token计算对7B模型大概每层几MB几十层加起来几百MB到1GB远小于KV Cache的数十GB。所以推理态的显存优先级排序是权重2P字节容易被低估但通常不是最危险的KV CacheL × 2 × hidden × precision × token数长上下文下最危险激活buffer按最大block决定总体小但可能在开大batch时被忽略依然可能触发OOM推理扩容时先看并发数和平均上下文长度如果平均上下文短如512权重占主导如果平均上下文长如8192KV Cache占主导。不同的主导方向决定了你是该加卡还是该换更省KV Cache的模型。4. 常见问题与排查技巧实录4.1 训练OOM了怎么快速判断卡在哪如果训练时CUDA out of memory先别急着把batch size减半先按如下顺序排查第一步用nvidia-smi看显存占用曲线。如果显存用量在一个step内出现尖峰后立即回落多半是激活峰值过高如果曲线平稳爬升后长期维持高位可能是优化器状态或梯度累计过多。第二步把batch size降到1seq len降到最短。如果OOM消失说明激活占大头如果依然OOM可能是权重优化器状态本身超了。第三步按前面的公式手算静态需求6P 4P 8P对比卡的显存。如果静态就超了batch size怎么调都白搭。4.2 推理服务并发上不去KV Cache 严重碎片化怎么办实测过一个场景8个并发请求上下文各有长短显存占用曲线看着还有空余但新请求不断OOM。原因就是KV Cache被零散请求切得七零八落无法分配连续块给新请求。解决方案就是类似虚拟内存的按块管理机制PagedAttention的实现思路把缓存切成小页按需映射。在实现层面这意味着你不能简单预分配“最大上下文 × 最大并发”的显存而要先弄清服务的平均并发、平均上下文、最大上下文这几个参数。KV Cache占比高时优先考虑给不同优先级请求不同长度的上下文限制。4.3 手算显存和实际占用差很多除了精度和GQA带来的偏差最常见的差异来源是“推理框架额外开销忘了算”。比如某些框架会为每个请求预留额外的激活buffer或者把中间结果保留用于日志分析。建议你拿到一个框架后先跑一个极小参数实验batch1seq1观察实际占用和公式的差把这个差当作该框架的“常数项”后续所有估算都加上它再考虑。4.4 我自己的避坑清单别在推理时用FP32权重除非精度需求极端否则内存直接翻倍到4P收益却很小。KV Cache量化不是默认打开的很多框架需要显式配置开关并配合校准数据集。梯度检查点确实能省激活但会多算一遍前向训练时长可能增加30%左右。省显存和加速往往是矛盾的要在资源约束下找平衡点。算FLOPs时记得把embedding的计算量也考虑进去。虽然比例可能不高但在小模型中可占5%~10%白板手算时别漏。最后分享一个实用小技巧每次拿到新模型先用自动化脚本输出四个数——参数量、每层hidden、层数、精度然后套两个公式一个是训练账本一个是推理账本存成表格贴在工位旁边。后续遇到OOM、并发瓶颈、升级硬件都可以先看这张表再动手改参数能省掉一大半瞎试错的时间。这套手算方法我用了很久从一开始的“算着算着算错”到后来能在白板上十分钟推完一个7B模型的完整显存和FLOPs分布过程中最有价值的一个体会是计算的目的不是得到一个精确到字节的数字而是得到一个可解释的、能在方向上指导决策的判断。真实世界中显存管理没有银弹但当你手里有一本清晰的账本时问题就已经解决了一大半。