显存到底被谁吃掉了

训练一个模型，显存远不止装参数本身。以常见混合精度训练为例，主要有四块：参数、梯度、优化器状态、前向保存的激活值。其中优化器状态常是最大头——例如 Adam 要为每个参数保存一阶矩、二阶矩，再叠加主权重，参数量乘以若干倍才是真实占用。

为什么优化器状态反而是大头？Adam 不只是存一份权重，它还要为每个参数维护两份「动量」（一阶矩、二阶矩）来估计自适应步长，再加一份高精度的主权重防止精度丢失。粗算下来，一个参数在 Adam 混合精度训练下要背十多字节，是权重本身的好几倍。这常常让初学者意外：他们以为「我模型才几十亿参数，一张卡不是能装几十亿吗」，却忘了梯度、动量、主权重全都要同步驻留显存。激活值则随批大小和序列长度暴涨——一句话几千个 token、几百层，中间结果在反向传播前都得留着。四块加起来，就是为什么训练比单纯「推理加载模型」吃显存得多。

三种并行：沿不同维度切

实际大模型训练几乎总是 3D 并行：节点内用张量并行、节点间用流水线并行、整体再叠加数据并行，在显存、通信和负载之间找平衡。

ZeRO：在数据并行内部再分片

普通数据并行每张卡都冗余存一份优化器状态、梯度甚至参数，浪费严重。ZeRO 的思路是把这些状态在各卡间分片，需要时再通过通信取回，分阶段（ZeRO-1 切优化器状态、-2 再切梯度、-3 连参数也切）逐级省显存，用可控的通信换显存，是当前训练框架的主流基座。

三种并行到底各切什么，可以用「怎么切一个蛋糕」来记。数据并行是把整蛋糕复制几份、每卡吃不同的那批客人（不同数据），最后汇总意见——简单，但每卡都得有整个蛋糕。张量并行是把一层的那把刀（权重矩阵）横着或竖着切成几份，每份卡只算一块，中途要频繁递来递去，所以适合放在互相连得极快的单机内。流水线并行是把蛋糕竖着切成一段段，第 0 层在 A 卡、第 1 层在 B 卡，数据像流水线一样一站站传过去。它们的区别本质上是在三个轴上做选择：切数据、切单层权重、切层序。现实里没人只用一种，而是 3D 组合——节点内用张量并行、节点间用流水线、整体再叠数据并行，在显存装不装得下、通信卡不卡、负载均不均之间反复调。

省显存与稳训练的常用手段

这几个手段有个共同逻辑：拿一样换一样。混合精度用「数值精度稍降」换显存和速度——BF16 比 FP16 数值范围更宽，所以训练更不容易溢出。梯度累积用「多算几批再更新一次」换「我显存小也能模拟大 batch 的稳定梯度」，本质上是把一次大更新拆成几步攒起来再一次性走。激活检查点最典型：前向时故意不保存全部中间结果，反向传播用到时再重新算一遍，用「多花一次前向的计算」换「省下大块激活显存」——显存实在不够时这是救命稻草，代价是慢一截。梯度裁剪则是给偶尔冒出来的超大梯度设上限，防止某一步把权重炸飞。没有哪个手段是白送的，工程师的工作就是在显存、速度、稳定性之间做对自己场景最合适的交换。