第 81 题：如何估算大模型训练所需的显存？公式推导？
题目
如何估算大模型训练所需的显存？公式推导？
完整讲解
一、显存组成
训练时显存主要来自：模型参数 (P)；梯度（与参数量同量级）；optimizer state（如 Adam 为 (2P) 的 fp32 动量+方差）；激活（与 batch、序列长、层数相关）。推理只需参数（及少量 KV cache 等）。
二、参数量与优化器
- 参数：fp16/bf16 为 (2P) 字节，fp32 为 (4P)。
- 梯度：通常与参数同精度，(2P) 或 (4P)。
- Adam state：两份 fp32，(8P)；若用 8bit optimizer 或 fused kernel 可减少。
- 总计（不含激活）：约 (4P + 8P = 12P)（fp16 参数+梯度 + Adam fp32）；ZeRO 分片后每卡除以 DP，再考虑 TP/PP 的切分。
三、激活的估算
- Transformer：每层激活约 (batch \times seq \times hidden \times (常数))，attention 还有 (batch \times head \times seq^2)；总激活约 (L \times (batch \cdot seq \cdot H + batch \cdot seq^2 \cdot H/L))，再乘 2（fp16）或 4（fp32）。用 checkpoint 可只存部分层，峰值约 (1/\sqrt{L}) 或按存点数量比例下降。
- 公式：显存 ≈ 参数+梯度+optimizer + 激活；激活 ≈ (L \cdot batch \cdot seq \cdot hidden \cdot C)（C 为常数与精度相关），或查表/经验系数。估算时先算 12P，再加激活项；激活项用 batch×seq×hidden×层数×系数 粗算。
面试要点
- 训练显存 ≈ 参数 + 梯度 + optimizer state + 激活；Adam 下约 12P + 激活。
- 激活与 batch、seq、L、hidden 相关；checkpoint 可显著降激活峰值。
- ZeRO/TP/PP 分片后每卡除以相应并行度，再按激活公式加总。
记忆要点
- 训练 ≈ 12P（fp16+Adam）+ 激活；激活 ∝ batch×seq×L×hidden。
- Checkpoint 降激活；ZeRO/TP/PP 分片降每卡参数与 state。
- 先算 12P 与激活量级，再按并行度除。