训练显存别再写成 20Ψ:把 Adam 状态与 ZeRO 分片分别记账
发布时间:2026/10/25 5:45:48 锦皓数字建站

一个训练方案写着“模型参数量乘 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 分布式训练、通信或数值回归。账本的作用是明确容量下界剩下的空间和精度风险仍要用目标训练栈验证。
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。