大模型显存计算法则——训练与推理全景图

本文目标：面试官问"一个 70B 模型 fp16 训练要多少显存？为什么混合精度训练 Adam 优化器要 16 bytes/参数？ZeRO-3 之后能降到多少？推理 KV cache 占多大？"——你能张口就算、并且讲清每个数的来源。

全文以**“算清训练显存 → 算清推理显存 → 进阶优化对显存的影响 → 面试速算题”**为主线。

---

一、前置：精度与字节换算（所有计算的基石）

精度	1 个参数占多少字节	常见用途
FP64 / double	8 B	几乎不用
FP32	4 B	优化器状态、master weight、科学计算
TF32	4 B（存储同 fp32，计算用 19 bit 尾数）	A100 训练加速，存储仍 4B
BF16 / FP16	2 B	训练前向/反向权重、推理权重
FP8 (E4M3/E5M2)	1 B	H100+ 训练/推理加速
INT8	1 B	推理量化、部分训练激活量化
INT4	0.5 B	推理量化（W4A16、GPTQ/AWQ）
INT3 / 2bit	0.375 / 0.25 B	极限量化推理

记住几个换算常数（后面反复用）：

• $1\text{B 参数} \times 2\text{B} = 2\text{GB}$（fp16/bf16 存权重）
• $1\text{B 参数} \times 4\text{B} = 4\text{GB}$（fp32 存权重或优化器状态）
• $1\text{B 参数} \times 16\text{B} = 16\text{GB}$（混合精度 + AdamW，每参数 16 bytes，下文详述）

面试金句："显存 = 参数量 × 每参数字节数 × 复制份数。训练显存主要由"权重 + 梯度 + 优化器状态 + 激活"四块构成，推理显存主要由"权重 + KV cache（+ 激活）“构成。”

---

二、模型参数量怎么算

2.1 Transformer 解码器层（LLM 主体）

一个标准 L 个层的 decoder-only Transformer，隐藏维 $d$，词表 $V$，注意力头数 $h$（每头维 $d_h=d/h$），MoE 暂不考虑，主要参数：

1. Embedding：$V \times d$（embedding 表；有时与 lm_head 共享权重，不共享则 ×2）。
2. 每层 Attention：
   • $W_q, W_k, W_v$：各 $d \times d$，共 $3d^2$。
   • $W_o$：$d \times d$。
   • 小项：RoPE 无参数；若用 bias，加 $4d$，可忽略。
   • 每层 attention 合计 $\approx 4d^2$。
3. 每层 MLP（标准 4 倍扩展，GatedSwiGLU 是 3 个矩阵）：
   • 标准 FFN：$d \to 4d \to d$，参数 $d\cdot4d + 4d\cdot d = 8d^2$。
   • GatedSwiGLU（Llama 系）：$W_{gate}, W_{up}$ 各 $d\times d_{ff}$，$W_{down}$ $d_{ff}\times d$，共 $3 \cdot d\cdot d_{ff}$。Llama 的 $d_{ff} \approx \frac{8}{3}d$ 取整到 64/128 倍数，故每层 MLP $\approx 8d^2$（与标准 FFN 量级一致，因为三个矩阵 × 较小宽度）。
   • 每层 MLP 合计 $\approx 8d^2$。
4. RMSNorm/LayerNorm：每层 2 个 norm，各 $d$，共 $\approx 2d$，可忽略。
5. lm_head：$d \times V$（若与 embedding 不共享）。

每层（不含 embedding/lm_head）合计：

$$\text{每层} \approx 4d^2 + 8d^2 = 12d^2.$$

总参数量近似公式：

$$\boxed{P \approx L \cdot 12d^2 + 2Vd\quad(\text{embedding 与 lm\_head 合计} 2Vd).}$$

经验：$Vd$ 相对 $L\cdot12d^2$ 通常很小（如 Llama-7B: $L=32, d=4096, V=32000$ → $12\cdot32\cdot4096^2 \approx 6.4\text{B}$，$2Vd\approx 0.26\text{B}$），所以粗算 $P \approx 12Ld^2$。

面试速算记忆点：一个现代 LLM，每层参数量 ≈ $12d^2$，再乘层数 $L$。见到一个新模型，看它的 $d$ 和 $L$ 立刻能估算参数量。

2.2 几个真实模型对一下

模型	$L$	$d$	$V$	$P$（公开值）	公式估算 $12Ld^2+2Vd$
Llama-7B	32	4096	32000	6.7B	6.4B+0.26B≈6.7B ✓
Llama-13B	40	5120	32000	13B	12.6B ✓
Llama-70B	80	8192	32000	65B	64B ✓
GPT-3 175B	96	12288	50257	175B	174B ✓

公式非常准，记住 $12Ld^2$。

2.3 MoE 的修正

MoE 层中 FFN 替成 $E$ 个 expert，每个 expert 是一个独立 FFN（$8d^2$ 量级），但每 token 只激活 $K$ 个。参数量按全部 $E$ 个 expert 算（都在显存里）：

$$P_{\text{MoE}} \approx L\cdot(4d^2 + 8d^2_{\text{shared}} + E\cdot 8d^2) + 2Vd.$$

而激活参数量（决定算力 FLOPs）只按 $K$ 个算。面试注意区分"总参数量"和"激活参数量"——MoE 显存按总参数算，算力按激活参数算。

---

三、训练显存：四块构成 + 混合精度 AdamW 的 16B/参数

训练显存由四部分组成：

$$M_{\text{train}} = M_W + M_G + M_O + M_A.$$

分别是 权重 / 梯度 / 优化器状态 / 激活。下面逐块拆。

3.1 权重 $M_W$

• 纯 fp32 训练：$M_W = 4P$ bytes（$P$ 参数个数，下同）。
• 混合精度训练（主流）：前向反向用 fp16/bf16，但保留一份 fp32 的 master weight 用于优化器更新（防止 fp16 下溢/精度丢失）。
  • fp16 权重：$2P$
  • fp32 master 权重：$4P$
  • 合计 $M_W = 6P$ bytes。

3.2 梯度 $M_G$

• 梯度通常和前向反向用同精度 fp16：$M_G = 2P$ bytes。
• （若 fp32 训练则 $4P$，但现代训练基本混合精度。）

3.3 优化器状态 $M_O$（面试最常算这块）

主流是 AdamW（带 momentum + variance 两个一阶矩，且存 fp32）：

• momentum $m$：fp32，$4P$
• variance $v$：fp32，$4P$
• （master weight 已算在 $M_W$ 里，不重复）
• 合计 $M_O = 8P$ bytes。

其他优化器：

• SGD with momentum：只有 1 个状态 $4P$。
• Adam（非混合精度）：权重本身 fp32 算在 $M_W$，状态 $m,v$ 各 fp32 共 $8P$，故纯 fp32 Adam 总计 $W+G+O = 4+4+8 = 16P$。等等——这正是下面"16"的由来之一。

3.4 混合精度 + AdamW 的总账：16 bytes/参数

把上面三块（不含激活）加起来（混合精度训练 + AdamW）：

项	每参数字节数
fp16 权重	2
fp32 master 权重	4
fp16 梯度	2
fp32 momentum	4
fp32 variance	4
合计	16

$$\boxed{M_W + M_G + M_O = 16P\ \text{bytes}.}$$

这是面试最常考的常数：混合精度 + AdamW，每参数 16 bytes。

快速算例（70B 模型，单机存不下要靠 ZeRO）：

$$70\text{B} \times 16\text{B} = 1120\text{GB} \approx 1.09\text{TB}.$$

即仅"权重+梯度+优化器"就要约 1.1 TB——一张 H100 80GB 远不够，所以必须 ZeRO / 模型并行分摊。

对比纯 fp32 + Adam（无混合精度）：

$$4(\text{W})+4(\text{G})+8(\text{O}) = 16P\ \text{bytes}.$$

诶巧合也是 16？注意来源不同：fp32 是 $W=4,G=4,O=8$；混合精度是 $W=6,G=2,O=8$。两者字节数都等于 16，所以常被混为一谈，但拆解不同，面试要能说清来源。

8-bit Adam（bitsandbytes）把 $m,v$ 量化到 8bit：$M_O = 2P$，总账降到 $6+2+2 = 10P$ bytes，省优化器显存。

3.5 激活 $M_A$（最容易被忽略、也最灵活）

激活显存 = 反向传播要重用的中间结果。每层激活占用（标准 attention + MLP，seq 长 $s$，batch $b$，$d$ 隐藏维，$L$ 层）：

粗略估计每层激活：

• attention 部分：注意力分数 $b\cdot h\cdot s^2$（每头一个 $s\times s$ 矩阵）+ QKV 输出 $b\cdot s\cdot 3d$ 等。
• MLP 部分：中间激活 $b\cdot s\cdot d_{ff}$，$d_{ff}\approx 4d$（SwiGLU 约为 $\frac{8}{3}d$，但中间更宽）。

经验公式（activation checkpointing 关闭时，每层 fp16）：

$$M_{A,\text{层}} \approx b \cdot s \cdot d \cdot (\text{几十的常数}),$$

常用近似 $\approx 10\sim 34 \cdot b\cdot s\cdot d$ 字节（不同文献常数不同，数量级正确即可）。关键是要会算激活 vs 权重谁占主导。

激活的工程控制手段——Activation Checkpointing（梯度检查点/重计算）：

• 不开：存全部中间激活，反向直接用。激活随 $s,b,L$ 线性增长。
• 开（标准）：只存每层输入，反向时重算前向，以算力换显存。激活降到约 $\sqrt{L}$ 层级 / 或每层只留 $b\cdot s\cdot d$ 量级。
• 开 + selective（如 FlashAttention 已自带）：只对耗显存大的注意力重计算，性价比最高。

面试要点：长上下文（$s$ 大）时激活显存往往超过权重显存，必须靠 activation checkpointing + FlashAttention（后者把 $O(s^2)$ 的注意力矩阵显存降到 $O(s)$）。

3.6 训练显存速算总表（混合精度 + AdamW）

场景	公式	70B 例子
权重+梯度+优化器	$16P$	1120 GB
+ 激活（无重计算）	$+M_A$（随 $s,b,L$）	+几十~上百 GB
+ 激活（有重计算）	$+M_A/\sqrt{L}$ 量级	+几~几十 GB

还要加：CUDA context、临时 buffer、通信 buffer（NCCL all-reduce 的 in-place 通常不额外占大块，但 reduce-scatter/ZeRO 分片有临时区）、框架开销。工程上预留 10%~20% 余量。

---

四、推理显存：权重 + KV cache（+ 少量激活）

推理没有梯度、没有优化器，显存干净：

$$M_{\text{infer}} = M_W + M_{\text{KV}} + M_{A,\text{小}}.$$

4.1 权重 $M_W$

• fp16/bf16 推理：$2P$。70B → 140 GB（一张 H100 80GB 不够，要 2 张 TP，或量化）。
• INT8：$P$。70B → 70 GB。
• INT4：$0.5P$。70B → 35 GB（单卡可跑）。

4.2 KV cache：推理显存第二大头，面试必算

自回归推理时，每生成一个 token 要访问之前所有层的 K、V。把它们缓存下来避免重算——这就是 KV cache。

单 token、单层 的 KV cache（fp16）：

• K 和 V 各占 $d$ 维（即 $d_h \times h = d$），共 $2d$ 个 fp16 = $4d$ bytes。
• 若用 GQA/MQA（grouped/multi query attention），K/V 的头数 $h_{kv} < h$，每层 KV cache = $2 \cdot d_h \cdot h_{kv} \cdot 2 = 4 d_h h_{kv}$。GQA 是省 KV cache 的关键。

全序列 KV cache（$L$ 层，$s$ 个 token，batch $b$，fp16）：

$$\boxed{M_{\text{KV}} = 2 \cdot L \cdot d \cdot s \cdot b \cdot 2\ \text{bytes} = 4Ldsb\ \text{bytes}\quad(\text{MHA}).}$$

（因子：2 = K+V；2 bytes = fp16。GQA 把 $d$ 换成 $d_h h_{kv}$。）

速算（Llama-70B, MHA, fp16, b=1, s=4096）：

• $L=80, d=8192$ → $M_{\text{KV}} = 4 \times 80 \times 8192 \times 4096 \approx 10.7\text{GB}$（单条 b=1 的 4k 上下文）。
• 若 batch=32 → 约 340 GB——KV cache 比权重还大，这就是为什么长上下文/大 batch 推理显存爆炸。

GQA 的价值：Llama-70B 用 GQA $h_{kv}=8$（$h=64$），KV cache 直接降到 1/8，约 1.3 GB（b=1, s=4k）。面试要点：现代模型用 GQA/MQA 大幅省 KV cache。

4.3 KV cache 的工程优化

• PagedAttention（vLLM）：像 OS 虚拟内存一样把 KV cache 按固定 block（如 16 token）分页管理，消除显存碎片、支持 dynamic batching、显存利用率从 ~20% 提升到 ~95%。面试高频。
• 量化 KV cache：KV 存 fp8/int8 降到一半甚至更少。
• KV cache offload：溢出到 CPU/SSD（如 vLLM 的 prefix cache、Mooncake/KVCache 池化）。
• Prefix caching / 共享 KV：系统 prompt 的 KV 在多个请求间复用。
• Multi-LoRA serving：多份 adapter 共享 base 的 KV。

4.4 推理显存速算示例

Llama-70B, bf16, 单请求 b=1, s=4096, GQA($h_{kv}=8$)：

• 权重：$2 \times 70\text{B} = 140\text{GB}$。
• KV：$4 \cdot 80 \cdot (8192/64)\cdot 8 \cdot 4096 / 2$… 简化用 GQA：$\approx 1.3\text{GB}$。
• 合计 $\approx 141\text{GB}$ → 单卡 80GB 不够，2 张 TP-2 切刚好。

如果 b=32 长 32k 上下文：KV 直接上百 GB，必须分页+量化+多卡。

---

五、并行与显存分摊：ZeRO / TP / PP / LoRA 的显存账

5.1 ZeRO（DeepSpeed）三阶段——面试必算

ZeRO 把训练状态在 $N$ 张卡间分片（sharding），消除数据并行里的冗余复制：

阶段	分片对象	每卡显存（权重+梯+优化器，16P 基础上）	相对单卡
ZeRO-0	不分，纯 DP	$16P$	1×
ZeRO-1	优化器状态分片	$\frac{8P}{N} + 2P + 2P = \frac{8P}{N}+4P$	优化器省 N 倍
ZeRO-2	+梯度分片	$\frac{8P+2P}{N} + 2P = \frac{10P}{N}+2P$	再省梯度
ZeRO-3	+权重也分片	$\frac{16P}{N}$	全分片，每卡 $\frac{16P}{N}$

速算 70B、16P=1120GB、8 卡：

• ZeRO-2：$10P/N + 2P = 1120\cdot10/16/8 + 140 \approx 87.5 + 140 = 227.5\text{GB}$ → 单卡 80GB 仍不够。
• ZeRO-3：$16P/N = 1120/8 = 140\text{GB}$ → 单卡 80GB 还不够（要 16 卡，140/16=87.5G 勉强）或叠加 TP。

注意 ZeRO-3 还有激活分片（DeepSpeed/Ulysses），以及通信代价：ZeRO-3 前向/反向都要 all-gather 权重，通信量比 ZeRO-2 大。省显存的代价是通信，这是面试区分点。

5.2 张量并行 TP / 流水并行 PP

• TP（Megatron-LM）：把每层权重矩阵按列/行切到 N 卡，每卡存 $\frac{1}{N}$ 权重 → 权重显存除 N，激活也分。通信是每层 all-reduce（机内 NVLink 才划算）。
• PP：按层切，每卡存 $\frac{1}{P}$ 层的权重+激活+优化器，通信量小但有 bubble。显存近似按层均匀分摊。
• 3D 并行：TP（机内）× PP（跨机）× DP/ZeRO（全局）。

5.3 LoRA / QLoRA 显存（参数高效微调）

• LoRA：冻结 base 权重（仍占 $2P$），只训低秩 adapter $\Delta W = BA$，秩 $r$。
  • 每层 LoRA 参数 $\approx 2 \cdot d \cdot r$（一对 down/up）。
  • 训练显存 = base 权重 $2P$（fp16，不存 master）+ LoRA 参数的权重/梯度/Adam 状态（量级 $O(Ldr)$，远小于 $P$）+ 激活。
  • 关键：base 不更新，不需要 base 的梯度和优化器状态，所以从 16P 直降到约 $2P + \text{小}$。
• QLoRA：base 权重量化到 4bit NF4（$0.5P$），加上 LoRA 训练 → 70B 可在单张 80GB 卡微调（$0.5\cdot70=35\text{GB}$ base + LoRA + 激活）。

面试金句：“LoRA 显存省的不是参数量本身，是省了梯度和优化器状态——冻结的 base 权重不需要梯度、不需要 Adam 的 $m,v$，把 16P 砍到约 2P（甚至 QLoRA 的 0.5P）。”

---

六、一张总图速查

训练（混合精度 + AdamW）                推理
─────────────────────────              ───────────
权重 fp16 ............ 2P              权重 fp16/int8/int4: 2P / P / 0.5P
master fp32 ......... 4P              KV cache: 4·L·d·s·b (MHA)  ← GQA可除 h/h_kv
梯度 fp16 ........... 2P              激活: 小（仅当前层，FlashAttn O(s)）
momentum fp32 ....... 4P
variance fp32 ....... 4P
──────────────────────              
合计 = 16P  (+ 激活, 可重计算压缩)      合计 ≈ 2P + M_KV (+ 少量)

并行分摊:
  ZeRO-1: 优化器分片      → 省 8P 到 8P/N
  ZeRO-2: +梯度分片      → 省 10P 到 10P/N
  ZeRO-3: +权重分片      → 16P/N （通信代价大）
  TP: 权重按矩阵切到机内 N 卡
  LoRA: base 冻结 → 约 2P；QLoRA → 约 0.5P + 小

---

七、面试速算题（自己先算再看答案）

Q1：Llama-70B 混合精度 AdamW 训练，单卡要多少显存？8 卡 ZeRO-3 够吗？

权+梯+优化器 = $16 \times 70\text{B} = 1120\text{GB}$。单卡 80GB 远不够。ZeRO-3/8 卡 = $1120/8 = 140\text{GB/卡}$，仍 > 80GB，不够；要么 16 卡（$1120/16=70\text{GB}$，加激活略紧），要么叠加 TP 或用 8-bit Adam（降到 $10P$ → $700\text{GB}/8\approx 87.5\text{G}$ 仍紧）。

Q2：为什么混合精度 + Adam 是"16 bytes/参数"？拆给我听。

fp16 权重 2 + fp32 master 4 + fp16 梯度 2 + fp32 momentum 4 + fp32 variance 4 = 16。本质是"一份 fp16 工作权重 + 一份 fp32 master 防精度退化 + 两份 fp32 Adam 矩"。

Q3：推理 70B bf16，单请求 4k 上下文，KV cache 多大？MHA vs GQA？

MHA: $4Ldsb = 4\cdot80\cdot8192\cdot4096\cdot1 \approx 10.7\text{GB}$。Llama-70B 用 GQA $h_{kv}=8$（$h=64$），降到 $\frac{8}{64}=\frac18$ → 约 1.3GB。权重 140GB + KV 1.3GB，TP-2 一张 80GB 刚好。

Q4：长上下文为什么激活显存会爆炸，怎么治？

激活 $\propto b\cdot s\cdot d$，注意力分数 $\propto b\cdot h\cdot s^2$，$s$ 大时 $s^2$ 的注意力矩阵显存爆炸。治法：FlashAttention 把注意力显存降到 $O(s)$；activation checkpointing 重算换显存；GQA 减小 KV。

Q5：LoRA 微调 70B 为什么单卡 80GB 能跑（QLoRA）？

QLoRA 把冻结的 base 存 4bit NF4 = $0.5\cdot70=35\text{GB}$，base 不更新所以无梯度/无优化器状态，只有 LoRA adapter（小）的权重+梯度+Adam + 激活。合计 < 80GB。要点是"冻结权重省掉了 14P 的梯度和优化器"。

Q6：ZeRO-2 和 ZeRO-3 的通信代价差别？

ZeRO-2 只分片优化器状态+梯度，权重每卡完整保留，前向/反向不需要 all-gather 权重，通信主要在梯度 reduce-scatter；ZeRO-3 权重也分片，每次前向、反向都要 all-gather 出完整权重再丢弃，通信量显著更大。所以 ZeRO-3 省显存但加通信，带宽不够时 ZeRO-2 更优。

Q7：一个 MoE 模型 671B（如 DeepSeek-V3），激活 37B，显存按哪个算？

显存按总参数 671B 算权重（bf16 约 1.34TB），因为它都在显存里；算力/激活 FLOPs 按激活 37B 算。MoE 的"省钱点"是算力省，不是显存省——显存仍是天量，必须 expert 并行（EP）把 expert 分到多卡。

Q8：训练 7B 模型，batch=1 seq=2048，激活显存大概多少？

每层 $\approx 34\cdot b\cdot s\cdot d$ bytes（fp16，经验常数）= $34\cdot1\cdot2048\cdot4096 \approx 0.28\text{GB/层}$，32 层 $\approx 9\text{GB}$（无重计算）。权重+梯+优化器 $16\cdot7\text{B}=112\text{GB}$。可见 7B 训练激活不是大头，但 70B/长上下文时激活反超。

---

八、口诀总结

参数量：$12Ld^2 + 2Vd$，MoE 总参数≠激活参数。
训练显存：权重+梯度+优化器+激活 = $16P + M_A$（混合精度+AdamW，每参数 16 字节）。
推理显存：$2P + M_{\text{KV}}$；KV cache = $4Ldsb$（MHA），GQA 按 $h_{kv}/h$ 缩。
省显存三板斧：ZeRO 分片（省 $P\to P/N$）、FlashAttn+checkpointing（省激活）、量化（$P\to0.5P$）、LoRA（冻结省梯度优化器）。
省显存必加通信：ZeRO-3、TP 都用通信换显存——工程永远是带宽与显存的权衡。

把 $16P$、$12Ld^2$、$4Ldsb$ 这三个常数和它们的来源吃透，大厂 ML 系统岗的显存题基本通杀。

---

参考

• Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, 2020.
• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022（激活显存精确推导）。
• Dettmers et al., QLoRA: Efficient Finetuning of Quantized LLMs, 2023.
• Kwon et al., Efficient Memory Management for LLM Serving with PagedAttention (vLLM), 2023.
• FlashAttention / GQA / DeepSpeed / Megatron-LM 文档。