训练显存别再写成 20Ψ：把 Adam 状态与 ZeRO 分片分别记账

一个训练方案写着“模型参数量乘 20 就是显存”，另一个写着“ZeRO-3 除以卡数就能装下”。两句话都省略了真正决定结果的假设。先拆状态，再谈切分，能少走很多容量规划的弯路。

16 字节来自哪里

在一种常见的混合精度 Adam 记账假设下，每个参数有 2 字节低精度权重、2 字节低精度梯度、4 字节 FP32 主权重、两个各 4 字节的矩状态，总共 16 字节。三个 FP32 状态合计是 12 字节，不是 16 字节。这也不是所有优化器、框架或训练配方的固定常数。

ZeRO 对应分片对象逐步扩展：优化器状态，随后梯度，随后参数。参考 ZeRO 论文时，应把状态分片和运行中的参数聚集区分开；参数在计算时仍可能临时聚集。只拿静态分片后的数字做 GPU 采购依据，会漏掉峰值。

图中只表达权重、梯度与两个矩状态的分片归属，未画 FP32 主权重；激活、通信缓冲和临时算子空间另计，不能从图上读取完整峰值。

先算状态下界

以下按上述 16 字节假设计算 7B 参数。它不下载模型，不创建张量，只检查分片算术。

GiB = 1024 ** 3 p, dp = 7_000_000_000, 8 cases = { "replicated": p * (2 + 2 + 12), "zero1": p * (2 + 2 + 12 / dp), "zero2": p * (2 + (2 + 12) / dp), "zero3": p * (2 + 2 + 12) / dp, } for stage, size in cases.items(): print(stage, f"{size/GiB:.2f} GiB/rank") assert cases["replicated"] == 112_000_000_000 assert cases["zero3"] * dp == cases["replicated"]

运行输出：

replicated 104.31 GiB/rank zero1 35.86 GiB/rank zero2 24.45 GiB/rank zero3 13.04 GiB/rank

即使状态估算降到了单卡容量之内，长序列激活仍可能占主导。微批数、激活重计算、参数预取窗口、通信 bucket 与 allocator 碎片都要单独测。增加数据并行卡数还会引入通信成本，不能假定时间也按比例缩短。

显存够了，数值也要对得上

10 月 14 日的训练与推理一致性研究讨论了相同权重在不同执行路径下的 log probability 差异。对 RL 训练来说，rollout 与训练前向路径的一致性是另一个需要检查的维度。低精度、算子选择和批布局的变化，都不能只靠“权重文件相同”判定等价。

实际评审可以保留一组固定输入，分别核对 logits、loss、梯度有限性和短程收敛趋势，再决定是否扩大集群。比较必须写清容差与采样设置，不能把一次脚本不报错称为精度无损。

本文本地只执行了 CPU 状态预算，未进行 PyTorch 分布式训练、通信或数值回归。账本的作用是明确容量下界，剩下的空间和精度风险仍要用目标训练栈验证。