激活显存估算

激活(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$(参数)+ 激活(重计算后),两者一起决定模型能否上卡。

推导激活显存公式 $\sum_l 2 \cdot s \cdot b \cdot h_l$: 解释每项来源(FP16 每元素 2 字节 × 每层输入激活 $h_l$ 个元素 × $s$ 个位置 × $b$ 个样本),并说明它为什么与参数量无关、随 batch×seq 线性增长。

数值例: Llama-2 70B($L=80$, $h=8192$), $s=4096$, $b=2$ → 每层 ≈ 134 MB、共 ≈ 10.7 GB;与静态 $16\Psi = 1120$ GB 并排讲清两本账各自随什么缩放(参数 vs batch×seq),即“互补”关系。

梯度检查点(重计算)的原理与开销: 每 $k$ 层存一个检查点、反向重跑前向,激活显存从 $O(L)$ 降到 $O(\sqrt{L})$,FLOPs 增加约 1/3;说明为什么是“省显存换算力”。

哪些张量最占激活显存(attention 的 $s \times s$ 分数矩阵、MLP 中间层 $4h$ 激活、dropout mask 等),选择性重计算如何只挑大头重算、把额外开销降到最低。

激活显存(训练)与 KV Cache(推理)的区别: 前者是前向中间输出、反向必须回读、随 $L \times s \times b$ 增长;后者是自回归生成的 K/V、顺序只读一遍、公式 $2 \cdot 2 \cdot L \cdot H \cdot d \cdot s \cdot b$;为什么两本账不能混算。