做百亿参数模型训练时，最头疼的不是loss不降，而是显存不够。很多同学拿着几张A100，满怀信心地开始训练，结果模型还没加载完就OOM了。这两年我在大模型训练上踩过不少坑，对显存占用这块算是摸得比较透了，这篇就认真梳理一下显存到底花在哪、又是怎么省下来的。
先说结论：大模型训练吃显存，主要不是吃模型本身，而是吃“训练过程产生的中间状态”。这句话理解了，显存优化的大方向就掌握了一半。具体来说，一块GPU在训练时要同时装下模型参数、梯度、优化器状态（Adam一阶二阶动量）、激活值（前向传播的中间结果）、临时缓冲区等，每一样都占地方。我拿一个具体的例子算过：一个7B参数规模的模型，用Adam优化器做FP16混合精度训练，光参数、梯度和优化器状态这三项就要占掉：
- 模型参数（FP16）：7B × 2字节 = 14GB
- 模型参数（FP32副本）：7B × 4字节 = 28GB（混合精度需要，下面细说）
- 梯度（FP16）：7B × 2字节 = 14GB
- Adam动量（FP32）：7B × 4字节 × 2 = 56GB
光这几项加起来就112GB了，还没算激活值。一张A100是80GB，单卡根本装不下。这也正是为什么需要各种优化手段，而且每招的着力点都不一样。下面我按“显存去哪了—逐层优化的招数—并行策略怎么选—避坑总结”这个顺序来拆。
1. 显存占用全拆解：训练时每一块显存都花在了哪里
1.1 静态显存与动态显存的区别
训练时的显存占用可以粗暴分成两类：静态占用和动态占用。
静态占用是模型结构决定的，加载之后就固定不动了。主要包括模型参数本身、梯度、优化器状态。这部分在初始化的时候就分配好，训练多久都不会变，变化曲线是平的一条线。动态占用则是训练过程中产生和释放的，典型的就是激活值、中间计算结果、临时变量，这部分在每一轮迭代中都会变化，跟批次大小、序列长度直接相关。你如果观察训练时的显存曲线，会发现它在一条基线上上下波动，这条基线就是静态占用，波峰就是动态占用。
很多人在排查OOM的时候只看模型大小，这是个误区。模型参数占的不是大头，优化器状态和激活值才是真正的显存杀手。特别是激活值，它和批次大小成正比——你把batch size翻一倍，激活值差不多也翻一倍，而参数和优化器状态不会变。
1.2 混合精度训练为什么会多出那么多显存
混合精度训练是现在的标配，但它是个典型的“用显存换稳定”的方案。FP16能省一半显存，但问题是FP16的精度不够：梯度在反向传播中容易溢出（超出FP16能表示的范围），导致训练不稳定。所以业界通用的做法是“FP16存储、FP32累积”：模型参数和梯度用FP16存，前向和反向计算用FP16算，但梯度更新这一步必须用FP32来做。这就要额外维护一份FP32的模型参数副本，还得专门给梯度更新准备一个FP32的“主权重”。
BF16也是一种选择，它跟FP16相比指数位更宽，数值范围大得多，训练稳定性更好，现在很多主流框架默认用BF16而不是FP16。但BF16的尾数位少，精度上还是有取舍，某些场景下需要配合Loss Scaling（损失缩放）技术。
1.3 激活值：最容易被忽视的大头
激活值是前向传播过程中每一层的中间输出。反向传播算梯度的时候，需要用到前向的激活值，所以不能算了就扔，得留着。这就是动态显存的主要来源。
一个很直观的计算方式：激活值占用 ≈ 批次大小 × 序列长度 × 隐藏层维度 × 层数 × 字节数。假设一个7B模型有32层、隐藏层维度4096，序列长度2048，batch size为1，每个值用FP16存（2字节），那光激活值这一项就是 1×2048×4096×32×2 = 512MB。看起来还能接受？但batch size变成32就是16GB了，再叠加上多卡并行时需要的通信缓冲，显存很快就兜不住。
这里有个从业者不太注意的细节：激活值不仅在前向时占显存，在反向传播的过程中还要继续占用，因为反向计算每一层的梯度时要逐层读取对应的激活值。这让激活值的生命周期比想象中长得多，也让重计算（Recomputation，下面会细说）这类手段有了发挥空间。
2. 主流显存优化手段逐个拆解：原理、代价与适用场景
2.1 梯度累积：不省显存，但能让你跑起来
先纠正一个普遍的误解：梯度累积并不能降低显存占用。它的作用是“用小batch size模拟大batch size”，让那些显存放不下大batch的训练也能跑大batch的等效效果。
原理很简单：正常的训练是攒够一个batch的梯度就做一次参数更新。梯度累积则是把N个小batch的梯度累加起来，攒够N个再做一次更新。这样每个step的batch size可以很小，显存压力自然小。但注意，梯度累积没有省掉任何计算量，只是改变了更新的频率。同时它会增加训练时间——本来N步就能完成的更新，现在要N×N步。
实际操作中我对梯度累积的建议是：如果单卡能跑batch size 4，你需要等效batch size 32，那就设梯度累积步数为8，同时开加速框架的梯度同步优化（比如通信与计算重叠）。还有一点，batch size太小时BatchNorm类的层会不稳定，Transformer架构影响不大，但卷积类的模型要谨慎。
2.2 混合精度深入细节：不止是“开了开关”
混合精度看起来就是改一个配置项，但里面的门道不少。最核心的机制是Loss Scaling（损失缩放）：上面提到FP16梯度容易溢出，一个解决方案是在前向计算时先把loss乘上一个比较大的系数（比如1024），让梯度在FP16范围内保持足够精度，反向传播完成后再把梯度除以这个系数，还原成真实梯度。如果发现梯度值过大导致溢出，训练框架会自动调整缩放系数，甚至会skip掉这一轮的参数更新。
经验上有几个值得注意的点：第一，越大的模型越适合用BF16，因为大模型的梯度分布更广，FP16更容易溢出；第二，开混合精度之后，BatchNorm层的统计量计算需要单独处理，PyTorch和主流框架都有对应的实现，但如果你手动实现训练循环，很容易在这块出bug；第三，Loss Scaling的值不是越大越好，过大会让梯度数值失真，反而不利于收敛，框架默认值一般够用，手调要小心。
2.3 激活重计算：用计算换显存的经典手段
激活重计算（Activation Recomputation，也叫梯度检查点）是“时间换空间”的代表作。思路是：前向传播的时候不保存全部激活值，只保存每层的输入（称为“检查点”）。反向传播需要某一层的激活值时，从检查点重新做一次前向计算，把激活值算出来再用。
简单来说，前向时本来要存N层的激活值，重计算只需要存N个检查点，反向时再花一次前向的时间把中间结果算回来。显存能从O(层数×序列长度×隐藏维度)降到O(序列长度×隐藏维度)，但训练时间会增加30%到50%左右。在我实际跑的模型上，这个代价通常是训练时间增加30%左右，换来的是显存占用减少50%以上，绝对划算。
实现上有两种粒度的重计算：一种是整个Transformer层级别的，每个层只保留输入，反向时重算该层全部激活值；另一种是更细粒度的，比如只重算Attention模块或者MLP模块。粒度越细，省显存越多，但重算开销也越大。多数实践里，按层粒度重计算是性价比最高的选择。
2.4 ZeRO优化器状态切分：数据并行下的显存革命
在数据并行（Data Parallelism，DP）模式下，每张卡都维护一份完整的模型参数、梯度和优化器状态，显存浪费非常明显。ZeRO的思想是：这些状态不一定每张卡都全量保存，可以切开，每张卡只持有自己的那部分，用到的时候通过通信把别人那部分取过来。
ZeRO有三个阶段：
- ZeRO-1：切分优化器状态。由所有卡分担Adam动量，理论上能节约4倍优化器显存（不切分时每卡都有完整动量）。
- ZeRO-2：在ZeRO-1基础上再切分梯度。每卡只负责自己那部分的梯度计算与更新。
- ZeRO-3：进一步切分模型参数。每卡只存一部分参数，用到其它参数时从别的卡广播拿。
用7B模型举例，ZeRO-1大约能省56GB的优化器状态（32卡场景），ZeRO-2额外省掉梯度的14GB，ZeRO-3再把模型参数的14GB省下来。所以ZeRO-3在卡数多时效果惊人，但代价是通信量显著上升，因为每次前向和反向都可能需要All-Gather参数和Reduce-Scatter梯度。
我实际测试过在32卡环境下，ZeRO-2是最稳的性价比选择。ZeRO-3虽然显存省得最多，但训练吞吐会明显下降，特别是在跨节点通信带宽不足的时候。如果你的卡数在8卡以内，ZeRO-2基本就够了。
2.5 参数卸载：把显存压力转移给CPU内存
ZeRO-Offload是把参数、优化器状态、梯度等挪到CPU内存，GPU只保留计算时的必要数据。这种做法能突破GPU显存的物理上限，代价是CPU和GPU之间的PCIe传输会成为瓶颈。
从实操来看，Offload是个“能跑但不快”的方案。在单个GPU且模型特别大的时候（比如单卡训13B模型），Offload能跑起来，但训练速度可能只有正常的十分之一。所以它适合“内存够、GPU差、只求能跑”的场景，比如在个人工作站上做小规模微调，不适合正经的大规模训练。具体配置一句话：能不用就不用，用了就做好一次训练跑好几天的心理准备。
3. 并行策略对显存的影响：怎么选才不亏
3.1 各种并行方式对比
大模型分布式训练里有几种并行策略，它们对显存的影响不太一样：
数据并行（DP）是最原生的方式。每个GPU都有一份完整模型，喂不同的数据。梯度同步用All-Reduce。如果配合ZeRO，就演变成了上面说的分片方式。数据并行在模型参数能塞进单卡时最有效，吞吐率高，实现也简单。
张量并行（Tensor Parallelism，TP）是把一个层切成多份，每张卡计算其中一部分。比如把4096维的FFN中间层切成4份，4张卡各算1024维。这样模型单层参数分散到多卡，激活值也会相应变小，但每层的前向和反向都涉及卡间通信，通信量非常大，通常只适合单机内部使用（NVLink带宽够）。超过单机规模，通信瓶颈就会显著拖累效率。
流水线并行（Pipeline Parallelism，PP）是把模型按层切成多段，每张卡负责其中一段。每张卡的激活值和参数都只覆盖自己的那部分层，显存压力自然小。但流水线并行有个“气泡”问题——同一时间有些卡在等待前面的卡算完，利用率会下降。好点的调度策略能减小bubble比例，但无法完全消除。
序列并行（Sequence Parallelism）是从另一个维度切分：沿着序列长度方向并行。它对显存的削减主要是针对激活值和Attention的计算，因为序列长度越长，激活值占用越大，而参数的量不变。
我自己的选型经验看，小规模的模型单卡或DP就够了；单卡装不下的中型模型（5B-13B）优先考虑TP+PP组合；超大模型（几十B以上）则要上3D并行（DP+TP+PP），配合ZeRO封装优化器状态。
3.2 异构显存场景下的工程技巧
实际工况下经常遇到“不完美的GPU池”：有8张A100，还有4张3090，甚至偶尔混着A800。这种情况下有几个实用技巧：
一是让显存大的卡扛更大的batch size，小卡用梯度累积补齐等效batch。实现上需要设置分布式训练的batch size参数为“单卡batch × 全局梯度累积”，并对各卡做差异化配置。
二是调整并行策略的切分比例：TP的切分会增加通信，同一个节点内多用TP，跨节点用PP和DP，可以减少跨节点小数据的频繁通信。
三是最简单也最容易被忽视的：用代码检查实际每张卡的显存占用，很多时候OOM是数据加载、中间变量、缓存没有清理导致，而不是模型本身。我见过一个案例，模型本身只占40GB，但因为dataloader里的worker进程没有释放显存，导致单卡直接OOM，排查了半天才发现是dataloader的pin_memory和分布式采样器没配好。
4. 从量化公式到实操：到底怎么算需要多少卡
4.1 手把手算一遍
掌握一个粗算公式能让你在做方案时快人一步。假设用Adam优化器和FP16混合精度训练L层Transformer模型（参数量为P），单卡训练大致的显存需求是：
- 模型参数（FP16）：2P字节
- 模型参数（FP32主权重）：4P字节
- 梯度（FP16）：2P字节
- Adam优化器状态：8P字节（4P×2，一阶和二阶）
- 激活值：跟batch size、序列长度、层数相关，粗略估算为 4 × batch_size × seq_len × hidden_size × L 字节
无优化时单卡显存基线通常约等于（2+4+2+8）= 16P + 激活值。这就是为什么7B模型最底限也要约112GB（不算激活值）的原因，80GB的卡单卡必然不够。
再算算ZeRO-2和ZeRO-3的数据。ZeRO-2在8卡环境下，优化器状态和梯度被切成8份，每卡只需要1/8，模型参数仍全量保留，那么每卡静态显存大约是：参数FP16 14GB + 主权重28GB + 梯度14GB×1/8 + 优化器状态56GB×1/8 = 50.75GB（不含激活）。这意味着80GB的卡大概能塞得下7B模型的训练，activation稍微注意点就能跑。而ZeRO-3继续切模型参数，每卡静态占用直接降到十几GB的量级，几乎不用担心显存，而是担心通信带宽和吞吐率。
4.2 一张表梳理各优化手段的取舍
我在这里放一个自查表，做方案的时候对照着选，效率会高很多。
| 优化手段 | 节省的显存对象 | 主要代价 | 适用场景 |
|---|---|---|---|
| 梯度累积 | 不直接省，降低峰值批大小 | 训练时间增加 | 显存不够但必须提高等效batch |
| 混合精度(FP16/BF16) | 参数与激活值的存储减半 | 精度损失，需Loss Scaling | 大模型训练标配 |
| 激活重计算 | 激活值（50%以上） | 30%-50%训练时间开销 | 激活值占用过高时 |
| ZeRO-1 | 优化器状态分摊 | 通信开销增加 | 多卡数据并行 |
| ZeRO-2 | 优化器状态+梯度分摊 | 通信开销增加 | 8卡以内最优解 |
| ZeRO-3 | 优化器+梯度+参数全分摊 | 通信开销显著增加 | 模型巨大且不追求速度 |
| Offload | 优化器/参数迁至CPU | 训练速度骤降 | 只求能跑不求速度 |
| TP/PP | 单卡模型参数减少 | 通信严重、代码复杂 | 超大规模模型训练 |
各手段不是互斥的，生产环境中几乎都是组合拳。最常见的高效组合就是：ZeRO-2 + 混合精度 + 激活重计算 + 梯度累积，四管齐下，7B级别的模型在8卡80GB环境下稳定训练没有压力。
5. 常见OOM案例与排查方法
5.1 快速定位OOM的排查思路
我遇到过很多次“莫名其妙”的OOM，最后发现都不是模型太大，而是踩了别的坑。遇到OOM先别慌，按顺序排查：
第一步，看日志和栈。PyTorch的OOM报错会给出发生OOM的代码位置，如果是在显存分配时爆的，通常是静态空间不够；如果是在迭代中间某个算子爆的，大概率是激活值太大。
第二步，用NVIDIA工具看显存曲线。观察每个step的显存使用趋势，如果显存曲线随step数增长而不释放，那就是显存泄漏。最常见的是张量没有被释放，累积在计算图里，或者是DataLoader创建了新张量但没有及时回收。
第三步，把batch size调小到1，看看能不能跑通。能跑通就是激活值或batch相关的问题，可以通过重计算或减小批次解决；跑不通就要看模型参数和优化器状态本身超了。把混合精度、ZeRO都打开后再试，通常能推断出瓶颈出在哪一部分。
5.2 一份可以直接抄的排查清单
| 异常现象 | 常见原因 | 解决方法 |
|---|---|---|
| 启动加载即OOM | 参数+优化器状态超过单卡上限 | 开ZeRO-2/3，或改Offload |
| 前几个step正常，跑久后OOM | 显存泄漏或数据累积 | 检查Tensor是否被引用，清空缓存；减小batch |
| 增大batch size后立刻OOM | 激活值占用过高 | 开激活重计算；梯度累积替代大batch |
| 开BF16后OOM | 主权重副本占额外显存 | 确认ZeRO优化器状态是否切分；检查是否有重复的FP32副本未清理 |
| 多卡训练时单卡显存异常偏高 | 并行配置不当 | 确认TP配置，单机内TP通信更优；检查数据加载是否有重复缓存 |
| 训练速度异常慢且显存稳定 | Offload或ZeRO-3通信瓶颈 | 换ZeRO-2；增加gradient checkpointing减少通信等待 |
| 指定卡不释放显存 | 框架缓存进程未释放 | 用工具查看残留进程，清理GPU缓存 |
5.3 关于重计算和通信重叠的几个观点
在实际训练中要想把吞吐率拉满，不能只盯着显存，还要看计算和通信能否重叠。比如ZeRO-2做梯度Reduce-Scatter的同时，学习的下一层回传计算已经在跑了，这就叫通信与计算重叠。框架层面现在很多已经把这类优化做成了默认配置项，但你要是自己写训练循环，就会出现“显存没爆但GPU利用率极低”的尴尬场景。
另外一个容易忽略的点：模型并行时，通信顺序和计算顺序强相关。TP通常先AllGather权重，再跑矩阵乘法，再做AllReduce，这中间的每一步如果都能和下一层的计算重叠，吞吐率就能接近线性扩展。相反，如果没有重叠，10张卡的训练可能只比单卡快2倍，这种情况我见过太多。
6. 从工程落地的角度聊聊怎么选方案
方案选择其实没有银弹，完全取决于你的具体场景。如果参数规模在1B以下，单卡甚至CPU都能微调，不用折腾多卡；如果参数在1B到7B之间，有两张以上80GB的卡，混合精度加ZeRO-2就够用了；如果到了13B到30B，必须上TP/PP的组合，配合激活重计算；再往上，比如30B甚至100B级别的训练，3D并行加ZeRO-3几乎是必然选择，还得考虑计算集群的拓扑结构。
参数规模与卡数量的对应关系，我给一个经验值：7B模型在8卡A100（80GB）上用ZeRO-2+混合精度+重计算，batch size可以开到8到16；13B模型同样的卡数，建议ZeRO-3或者TP+PP，batch size要降到4到8；30B模型在32卡环境下，TP=4、PP=4、DP=2的组合比较均衡。这些都只是起始参考点，实际要跑一两个step看显存和吞吐率再调。
最后分享一个我个人踩过多次坑之后的习惯：做任何大模型训练任务之前，先跑一个最小单元（batch size=1、单卡、短序列）来验证模型结构和前向反向流程，然后再逐步放大batch、序列长度、卡数量。这样做能帮你把“模型写错了”和“显存不够”这两个问题分开排查——要知道把模型代码bug误判成显存问题，可是浪费过我整整两天时间的。
显存优化的本质不是把每个字节都抠到极致，而是在你的硬件预算和训练效率之间找一个平衡点。希望这篇能帮你在配方案的时候少走一些弯路。
