激活显存估算

激活(activations)是前向传播中每层的中间输出,反向传播要回读它们计算梯度,所以训练时要么存、要么重算。与参数静态显存 $16\Psi$“只随参数量、与 batch/seq 无关”相反,激活显存正比于 批大小 $b$ × 序列长 $s$ × 隐藏维度 $h_l$。每层保存输入激活(FP16 每元素 2 字节)需 $2 \cdot s \cdot b \cdot h_l$ 字节,累计得激活显存 ≈ $\sum_{l=1}^{L} 2 \cdot s \cdot b \cdot h_l \approx 2 L s b h$ 字节;经验法则: 每层每 token 约 16 字节(≈ 每 token 存 8 个 FP16 元素的粗估口径,视隐藏维度而定),总激活 ≈ $16 L s b$ 字节,精确值用 $\sum_l 2 s b h_l$ 计算。数值例: Llama-2 70B($L=80$, $h=8192$), $s=4096$, $b=2$ → 每层 $2 \times 4096 \times 2 \times 8192 \approx 134$ MB,共 80 层 ≈ 10.7 GB——只有静态参数显存(1120 GB)的约 1%,但随 $s \times b$ 线性放大,长上下文 + 大批次时反客为主。重计算(Gradient Checkpointing): 不存每层激活,只每 $k$ 层存一个检查点,反向时重跑前向把中间激活算回来;激活显存从 $O(L)$ 降到 $O(\sqrt{L})$(取 $k \approx \sqrt{L}$),FLOPs 增加约 1/3——省显存换算力;选择性重计算只重算最占显存的张量(attention 的 $s \times s$ 分数矩阵、MLP 中间激活),把额外计算压到最低。

“70B 训练到底要多大显存”的完整账本必须叠加激活项;面试追问“推导激活显存公式”“梯度检查点原理与代价”“激活显存与 KV Cache 的区别”。

只算 $16\Psi$ 会严重低估长序列/大批次训练的真实显存(激活可反超参数显存);重计算用约 33% 的额外计算把激活显存从 $O(L)$ 压到 $O(\sqrt{L})$,与 ZeRO 参数切分互补后,显存预算 = $16\Psi/N_d$(参数)+ 激活(重计算后),两者一起决定模型能否上卡。

检验攻克程度：针对「激活显存估算」专属刷题排雷

做单选排雷题、推导选项机制，答错自动收录进专属错题本。