跑大模型训练的人多多少少都经历过这样的瞬间显卡显存监控明明还剩二三十GB任务却突然弹出一行“CUDA out of memory”推理服务昨晚还好好的今天并发一上来就开始批量拒绝请求。很多人第一反应是“加卡”加上卡之后又开始犯嘀咕——显存到底去哪儿了这一讲我们不讲玄学只讲账本。我会把大模型训练和推理阶段的显存占用拆成几大块让你能拿白板一步步手算训练侧有参数、梯度、优化器状态、激活值推理侧有权重、KV Cache、临时缓冲区。顺手再把FLOPs怎么在纸上估KV Cache容量怎么规划一起讲透。别被“第27讲”这个序号吓到其实核心问题就一句话每一份显存花在了哪里怎么用公式把它算出来。适合刚接触超大模型的工程师也适合想把推理成本算清楚的同学。有这本账在手你至少能做到三件事训练任务启动前提前判断单卡够不够、要不要用并行策略推理服务上线前直接算出来最大并发对K是几别人跟你说“这个模型只要几块卡就能跑”的时候你能第一时间分辨他说的是训练还是推理。1. 显存账本第一页参数、梯度与优化器状态的三方分立1.1 一个参数在训练时到底占几份很多人对显存的第一反应是“模型权重多大就占多大”。这个直觉在推理场景勉强成立但在训练场景完全不成立。训练时一个参数往往同时存在好几份副本先把这个基础算清楚项目每参数字节数说明模型权重2混合精度训练时前向和反向用的是BF16/FP16权重梯度2反向传播产出的梯度通常也按BF16/FP16保存优化器状态12AdamW需要FP32的master权重副本、一阶动量、二阶动量每项4字节相加之后一个参数在典型混合精度训练里大约占到2 2 12 16字节。这是个非常关键的常量建议直接背下来。为什么要存FP32的master权重因为BF16/FP16的精度不足以完成优化器更新。AdamW在更新参数时如果把很小的学习率和很低的精度放一起数值更新很容易被抹掉或溢出所以优化器里始终保存着一份FP32的“标准答案”每次更新在这份副本上完成再转回BF16供前向和反向使用。按16字节/参数计算一个70B模型训练阶段的显存底盘就是70e9 × 16 ≈ 1120GB。这意味着如果你想在单卡上直接把70B模型加载起来做微调哪怕不算激活值也需要十几张80GB显卡。如果有人跟你说“4张80GB卡就能装下70B模型跑训练”他大概率是把推理场景的权重占用和训练场景搞混了。1.2 ZeRO 到底切的是什么既然训练显存的大头是优化器状态并行策略的切入点就很清楚了。DeepSpeed里的ZeRO本质上是把“参数、梯度、优化器状态”这三类东西在不同GPU之间切分Stage 1切分优化器状态每张卡只持有自己负责的那份更新状态Stage 2再加上梯度的切分Stage 3连参数也一起切分每张卡只放自己需要的参数分片用一个仓库类比一间仓库放不下全部货物与其硬塞进去不如分到多个仓库谁要用货时再临时去取。代价是通信量变大但显存压力直线下降。实际做微调时还有个容易忽略的点很多人用Hugging Face生态里现成的框架默认帮你开了ZeRO或者FSDP于是对显存账本没有体感。一旦换成裸PyTorch手写训练循环瞬间就会撞上OOM。这不是代码写得不好而是你没有意识到“参数16字节/参”这本底账的存在。先把这个常数刻在脑子里后面的推理账本才有依据。2. 训练显存的“隐形巨头”激活值为什么比权重更占地方2.1 激活值是什么为什么不能随手删如果说参数、梯度、优化器状态是清晰摆在明面上的账目激活值就是藏在报表中下方的“其他项”而且它经常是训练场景里最大的一项。激活值指的是前向传播过程中每一层的中间张量。比如Attention里的Q、K、V矩阵注意力分数MLP中间层的临时结果。反向传播计算梯度时需要用到这些中间结果所以它们不能算完就扔。这就像装修施工现场的脚手架楼盖到一半脚手架拆了工人就没法上去刷墙。激活值的量级和什么相关它不直接依赖参数量大小而是和网络层数、hidden维度、序列长度、batch size强相关。一个7B模型可能激活占用比一个13B模型更高只要前者的序列长度和batch更大就会发生。2.2 一张白纸就能算激活值的估算公式GPT系列论文里给过一个被广泛引用的激活显存估算公式单位是字节激活内存 层数 × batch size × 序列长度 × hidden维度 × (34 5 × 序列长度 × 注意力头数 / hidden维度)这个公式是怎么来的括号里第一项是Self-Attention和MLP里各种固定形状的中间矩阵第二项是注意力分数矩阵它随序列长度线性增长。不需要精确到字节级能给的是一种数量级判断但已经足够做容量规划。拿一个常见的7B模型来算hidden 4096层数 32注意力头数 32序列长度 2048batch size 1。括号里 5 × 2048 × 32 / 4096 80所以总数为 32 × 1 × 2048 × 4096 × (34 80) ≈ 30.6GB。对比一下这个7B模型如果是BF16权重文件大小约14GB。一层的计算里仅仅保存激活值就要吃掉30GB出头是权重的两倍多。而且这里用的还是“单条样本、2048长度”的保守场景真要把batch size拉到8、把序列长度拉到4096激活可以冲到三百GB级别全部训练显存的大头就此易主。2.3 激活值省显存三板斧第一招是激活重计算也叫activation checkpointing。思路很直接前向传播时只保留每一层最边缘的少数关键中间结果反向传播要用到中间张量时现场重新做一次前向把它们算出来。代价是大约30%的额外计算量换来的是激活显存可以降到原来的十分之一甚至更低。常规大模型训练基本默认开启。第二招是梯度累积。它不减少全局batch size而是把一个大batch拆成多个micro-batch每个micro-batch只走一小段前向和反向梯度和优化器更新攒够了再统一执行。激活峰值只按单个micro-batch算显存压力大幅下降。第三招是流水线并行。把不同层放到不同GPU上每张卡只保留当前这截模型的激活值。它解决不了单层显存峰值但能有效削减单卡总量和ZeRO是互补关系。实际操作里的经验是先开重计算再把micro-batch压到能塞进卡里为止最后才考虑是否上流水线并行。为了一个本来可以靠重计算解决的激活爆炸问题直接拉出一套复杂的并行方案通常是过度设计。3. FLOPs 不算清楚集群规模就是拍脑袋3.1 FLOPs和FLOPS是两个东西很多人把FLOPs和FLOPS混着用这两个概念差得其实很远。FLOPs是浮点操作数是一个总量我们平时说“训练这个大模型需要多少计算量”用的就是FLOPs。FLOPS是Floating Point Operations Per Second是设备的吞吐能力。类比一下FLOPs是路程FLOPS是车速拿车速除以路程就知道要跑多久。显卡厂商给的“峰值TFLOPS”只是理想车速实际训练里因为通信、算子效率、显存搬运等原因用到的MFU模型算力利用率往往只有40%到60%。所以手算FLOPs时一定要给车速打个折扣否则你算出来的时间和实际时间会差一倍还多。3.2 一个token配一个参数为什么是6次浮点操作训练一个标准Transformer模型的总FLOPs业界有个特别好用的近似公式训练FLOPs ≈ 6 × 模型参数量 × 训练token总数6这个系数是有出处的。对每个参数而言前向传播大约需要2次浮点操作反向传播要分别对输入和权重求梯度工作量大约是前向的两倍也就是4次。前向加反向正好6次。这个系数对大多数自回归语言模型都成立足以支撑数量级层面的判断。如果你问的是推理阶段生成一个token那公式变成推理单token FLOPs ≈ 2 × 模型参数量不包含反向传播只有前向所以系数大幅下降。还要注意这套近似公式主要覆盖的是Transformer核心块embedding层和最后的输出层在超大模型里占比很小可以忽略。但如果你的模型用了超长序列Attention矩阵乘法的二次方开销会额外增加公式就只能当“地板价”来看了。3.3 手算案例70B在2T token上要烧多少卡假设要在一个70B模型上训练2T token。这个规模在业内很常见我们现场算一遍总FLOPs 6 × 700亿 × 2万亿 8.4 × 10²³。用1000张A100 80GB来算它的BF16峰值大约是312 TFLOPS也就是每秒3.12 × 10¹⁴次浮点操作。按45%的MFU打折每秒有效算力约1.4 × 10¹⁴。一张卡一天能提供的有效FLOPs约为1.2 × 10¹⁹1000张卡一天约1.2 × 10²²。8.4 × 10²³ ÷ 1.2 × 10²² ≈ 70天。所以你得到结论1000卡A100跑这样一个训练任务大约需要两个多月。如果换成峰值算力更高的卡或者MFU优化得好时间可以进一步压缩反之通信瓶颈和长序列会让天数明显变慢。这个手算过程看起来粗糙但它是一个非常好的“现实检查工具”。当有人抛给你一份训练方案你先把参数量、token数、卡数、单卡峰值算力对上就能立刻判断方案是不是在吹牛。做推理服务的容量规划也一样先算总量再谈优化才不会出现“显存看起来很大并发一开全废”的尴尬。4. 推理阶段显存格局巨变KV Cache如何登堂入室4.1 Prefill和Decode两种阶段的显存性格完全不同推理阶段可以分成两段看。Prefill阶段一次性吃下整个输入prompt把所有历史token的中间结果并行算出这个阶段计算密集峰值显存往往很高因为短时间内要同时保留很多激活矩阵。Decode阶段是逐token生成输出每一步只算最新的一个token计算量不大但每一步都要读取历史token的Key和Value所以显存占用会随输出长度持续累积。KV Cache就是这段累积的结果。它不是某个TOKEN的一次性计算结果而是为了后面所有决策都必须长期保存的“记忆”相当于每次写文章时要记住前面已经写过什么。有人想用类似训练里activation checkpointing的方式重算KV那在数学上可行但会把推理速度拖慢到不可接受所以业界几乎不这么干。4.2 KV Cache容量公式一行代码算出显存KV Cache的占用公式可以写成这样KV Cache字节数 2 × 层数 × batch size × 序列长度 × KV隐藏维度 × 每元素字节数这里的2是因为每层要分别保存Key和ValueKV隐藏维度在标准多头注意力里等于hidden size如果用了GQA等变体这个值会按组数缩小。用代码表达会更清楚比如用一个Python函数直接计算def kv_cache_bytes(layers, batch_size, seq_len, kv_hidden_dim, bytes_per_element2): # 2 表示 K 和 V 两份 return 2 * layers * batch_size * seq_len * kv_hidden_dim * bytes_per_element # 7B模型, MHA, hidden4096, layers32, 单请求, 8192上下文 print(kv_cache_bytes(32, 1, 8192, 4096, 2) / 1024**3, GB) # 输出约 4.0 GB看起来单请求的不算离谱。但注意这个公式里有一个乘数是batch size也就是并发请求数。同一个7B模型如果同时服务100个并发请求上下文都按8192算KV Cache就会膨胀到400GB级别。权重才不过14GBKV Cache反手甩开权重三十倍。这就是推理服务“并发一上来就OOM”的根源。极端情况下更夸张。一个70B模型hidden 8192层数 80如果上下文长度是128K单请求KV Cache就能到312GB左右比模型权重140GB还高出一大截。所以长上下文推理绝不是“把token塞进去”就完事KV Cache容量的增长速度会让人措手不及。4.3 MHA、GQA、MQA结构上的KV减负既然大头在KV那从模型结构上砍KV就最有效。标准多头注意力MHA每一对query头都有独立的Key和ValueKV隐藏维度和hidden完全一致。GQA做的是让一组KV头服务于多个query头组内query共享KVKV隐藏维度直接除以分组倍数。MQA更极端所有query头共用一组KVKV cache压到原来的层头数分之一。注意力类型KV头数KV隐藏维度相对KV CacheMHA query头数 hidden1倍GQA8组为例约query头数/8hidden/81/8MQA1head_dim1/query头数比如70B模型做128K长上下文MHA约312GB改成8组KV的GQAKV cache能降到40GB左右容量规划瞬间从“完全不可行”变成“可行但仍有压力”。业界后来的超大模型普遍采用GQA不是性能上一定更优而是在KV Cache这个显存约束面前结构已经先替工程省钱省到了。5. KV Cache 容量规划从单请求推演到整机并发5.1 容量规划的四步法我通常按照四个步骤来做推理服务的显存规划每一步都能落到具体数值第一步算权重占用。BF16/FP16按每参数2字节INT8量化按1字节INT4按0.5字节。例如80B模型BF16权重约160GB。第二步预留运行时和激活。哪怕是推理Prefill阶段也需要一块临时显存来放激活矩阵和算子中间结果加上框架本身的缓冲通常按模型权重的20%到30%预留。长序列请求会吃更多建议多留一点。第三步拿总显存减掉前两步剩下的就是KV Cache可用空间。这里要扣除的是设备实际情况比如单卡80GB但实际还有系统预留主动给“不可用显存”留出几位。第四步算出单请求KV Cache占用用第三步的余量除以它得到理论最大并发。单请求KV Cache的计算用前面那个公式就好。注意这里的序列长度应该按“平均prompt长度 平均输出长度”来估算不能只看当前的短样本否则上线后长请求一来就会出问题。5.2 实例8卡80GB服务一个80B模型听上去很抽象不如实际代入一组数字模型80BBF16权重160GB8张80GB卡总显存640GB。预留运行时和临时激活40GB剩余可用KV Cache约440GB。如果是MHA假设hidden8192层数80那么每个token的KV Cache占用是2×80×8192×2字节约2.5MB/token。平均每个请求的prompt加输出长度按2048 token算单请求占用约5GB。440GB除以5GB理论并发约88路。如果换成一个8组KV的GQA版本KV隐藏维度变成1024单token占用约0.31MB单请求约0.63GB那同样440GB就能支撑约700路左右。差距就是这么明显。但这里还有个工程现实同时跑88路时每路解码速度会受显存带宽和计算资源的限制显存够不够只是必要条件不是充分条件。真上线时还要配合压测观察解码延迟和排队情况再决定最高并发到底设多少。容量规划只能告诉你“显存天花板在哪”不能替代负载测试。5.3 量化、缓存复用和分页管理KV Cache最直接的优化手段是量化。FP16换成FP8占用直接减半再到INT8、INT4容量还能进一步压缩。代价是KV信息的精度损失尤其对长上下文和敏感任务需要把量化后的效果放到具体业务里验证。很多时候模型能力没明显下滑但也不能把这个过程当默认操作。另一个思路是前缀缓存。如果很多请求共享同一段系统提示词或指令前缀那这段前缀算出来的KV Cache完全可以复用不用每来一个请求就重算一遍。这在高并发QA场景里效果拔群相当于把重复计算整段删掉。分页式KV管理则是从内存碎片角度出发把KV Cache拆成固定大小的页来分配避免不同长度的请求造成大量碎片浪费。这个思路有点像操作系统的虚拟内存原理不复杂但工程收益巨大。目前很多开源推理框架都在朝这个方向演进规划显存时建议选择支持类似机制的方案能显著减少实际峰值占用。6. 最后落在实操看显存不能只看nvidia-smi流程上要算账先行6.1 用对工具分清“占用”和“已分配”的区别我见过太多人对着nvidia-smi的显存数值一顿分析其实那个数字反映的是驱动层面的显存占用而PyTorch这类框架自己维护着一个缓存分配器。框架用完一块显存后不会立刻还给驱动而是留在自己池子里复用。结果是nvidia-smi看起来占得很满但框架内部可能还有大量空闲块可以分配反过来也成立——进程申请了大块显存池驱动还没来得及显示任务就OOM了。所以在PyTorch里排查我更推荐用torch.cuda.reset_peak_memory_stats() # ... 跑一段训练或推理 ... print(torch.cuda.max_memory_allocated() / 1024**3)max_memory_allocated()统计的是框架实际分配的峰值远比nvidia-smi能反映真实逻辑。另外还可以用torch.cuda.memory_summary()看块级碎片情况判断是不是分配了很多无法复用的细碎片段。这是排查OOM时最容易走弯路的环节工具没选对方向就全歪了。6.2 三条反直觉的实操心得第一条训练场景里先算账的优先级永远高于先调代码。一个70B规模的模型什么优化都不做直接裸训练显存需求按1120GB起步这个数字一出来就知道单机方案没戏。先定并行策略再开任务能省掉一整晚的OOM调试。第二条推理和训练的显存结论不能互相套。训练讲“它需要一个庞大的优化器状态”推理完全用不到推理讲“KV Cache随并发增长”训练又没这个负担。每当我看到有人拿推理显存去反推训练可行性就知道他该回来看账本了。第三条FLOPs和显存必须一起看。开activation checkpointing能省显存但会额外增加30%左右的计算量KV Cache量化能省显存但可能影响生成长度条件下的效果。任何一个优化动作都是一笔“显存换计算”或“精度换容量”的交易先知道总数再谈交易才不会顾此失彼。我自己现在的习惯是接到一个项目先花半天时间拿白板把训练FLOPs、训练显存、推理KV Cache三张表各算一遍所有假设都落在白纸黑字上然后才去申请卡、选服务框架。这个习惯帮我避开过很多“看着能跑、一跑就崩”的坑。显存这件事看起来是硬件问题本质上是一本随时能翻出来的预算账账算明白了人和卡都轻松。