Transformer Math：FLOPs、参数、显存、KV Cache 精确算
Transformer Math：FLOPs、参数、显存、KV Cache 精确算
前情说明：技术深度 D2。《读 JAX Scaling Book》系列第 3 篇。
前两篇讲了"一个 kernel 该多快"（Roofline）和"分片矩阵乘怎么算"（Sharded Matmul）。这一篇解决最基础也最常被问倒的一类问题:
- 这个模型多少参数? 不用查表,自己算
- 训练要多少 FLOPs?
6 × 参数 × token这个公式从哪来- 一个 batch 要吃多少显存? 激活值怎么估
- KV Cache 多大? 能跑多少并发
- attention 什么时候开始比 FFN 更贵? 有个精确临界点
面试被问"Llama 70B 部署到 8 卡 H100 能撑多少并发"——这篇文章学完之后你可以在白板上当场算出来。我把 Llama 3 70B 和 DeepSeek V3 的完整数字都算好了,可以直接对照。
一、从数"点"开始:FLOPs 的基本单元
一切从最基础的操作出发。设向量 x, y ∈ [P],矩阵 A ∈ [N, P], B ∈ [P, M]:
操作 FLOPs 数据量
─────────────────────────────────────────
x · y (点积) 2P 2P
A x (矩阵向量) 2NP NP + P
A B (矩阵矩阵) 2NPM NP + PM
三个必须记住的:
- 点积 P 个元素 → 2P FLOPs（P 次加 + P 次乘）
- 矩阵乘 → 2NPM（每个输出元素是 N 维点积，共 PM 个输出）
1.1 一般化的 einsum 规则
对任意维张量收缩,FLOPs 有一个极简规则:
给定 C[GH, IJ, KL] · D[GH, MN, KL] (KL 是 contracting, GH 是 batching)
FLOPs = 2 × (所有维度乘积,其中 batching/contracting 维只算一次)
= 2 × G × H × I × J × M × N × K × L
术语澄清（这是原书讲得很清楚的一点）:
- contracting 维度：在两个输入里都有、输出里没有 → 被求和掉（矩阵乘的 inner dim）
- batching 维度：在两个输入里都有、输出里也保留 → 索引独立的子问题,不算在乘法里
einsum 视角:两个输入都有 + 输出也有 = batching;两个输入都有 + 输出没有 = contracting。
1.2 为什么 Transformer 偏爱 matmul
矩阵乘 N×N 的规模扩张:
N 增大 → FLOPs 按 N³ 增长
数据搬运按 N² 增长
─────────────────────
Intensity ∝ N
这是极其罕见的性质!
意味着矩阵越大越容易吃满硬件(回想 Roofline)
对比: element-wise 操作
FLOPs ∝ N, Bytes ∝ N → intensity 恒定 ~1
→ 永远 memory-bound,永远榨不干硬件
这就是 Transformer 一切设计围绕 matmul 转的根本原因——它是唯一"越大越高效"的算子。
二、6 × 参数 × token 这个公式从哪来
这是深度学习界最著名的经验公式。必须能自己推。
2.1 单个矩阵乘的反向传播
设网络里有一个 C = A · B,其中 A ∈ [N, P] 是输入激活,B ∈ [P, M] 是参数,C ∈ [N, M]。
Forward:
C = A · B
FLOPs = 2NPM
Backward (需要两个梯度):
dL/dB = A^T · (dL/dC) ← 对参数求导, [N,P]^T · [N,M] → [P,M]
FLOPs = 2NPM
dL/dA = (dL/dC) · B^T ← 对输入求导(传给前一层), [N,M] · [P,M]^T → [N,P]
FLOPs = 2NPM
─────────────────────────────────────────
总计: fwd 2NPM + bwd 4NPM = 6NPM
关键:PM 正是参数个数！所以:
单个矩阵乘训练 FLOPs = 6 × 参数量 × N
↑
N = 这一层的 token 数
推广到整个 Transformer:
Total FLOPs = 6 × 总参数量 × 总 token 数
这就是 6 × params × tokens 的完整推导。它不是一个经验数字,是从 fwd 1x + bwd 2x 严格推出来的。
2.2 三个倍数必须记住
┌──────────────────────────────────────────┐
│ Forward : 2 × params × tokens │
│ Backward : 4 × params × tokens │
│ 训练总计 : 6 × params × tokens │
│ │
│ 推论: 训练 FLOPs 永远是推理的 3 倍 │
└──────────────────────────────────────────┘
验证一下 Llama 3 70B:15T tokens 训练
FLOPs = 6 × 70.55e9 × 15e12 = 6.35e24 FLOPs
Meta 实际用了多少卡时?
如果用 H100 (989 TFLOPs, 假设 40% MFU):
6.35e24 / (989e12 × 0.4) = 1.6e7 s = 4460 小时·卡
→ 约 6400 张 H100 跑 16 天
2.3 这个公式什么时候失效
重要边界（面试常被追问）：
✓ 成立: 参数占主导的 dense Transformer, 短上下文
✗ 失效场景:
1. 长上下文 → attention 的 S² 项变大(见第四节)
2. MoE → 只有"激活参数"参与计算,不是总参数
3. 大量 element-wise 操作(罕见)
4. 参数量极小的模型(embedding 占比大)
三、完整 Transformer 的账本
现在逐块算。用标准 decoder 记号:
B = batch size (序列数)
T = 序列长度
D = d_model (hidden dim)
F = d_ff (FFN 中间维, 通常 4D 或 8D/3)
L = 层数
N = Q head 数
K = KV head 数 (MHA: K=N; GQA: K<N; MQA: K=1)
H = 每个 head 的维度 (通常 N·H = D)
V = vocab size
3.1 MLP 部分
LLaMA / DeepSeek 用 gated MLP（三个矩阵,不是两个）:
操作 训练 FLOPs 参数
────────────────────────────────────────────────────────────
A[B,T,D] · W_in1[D,F] 6BTDF DF
A[B,T,D] · W_in2[D,F] 6BTDF DF
SiLU(A_in1) * A_in2 (element-wise) ~O(BTF) 0
A[B,T,F] · W_out[F,D] 6BTDF DF
────────────────────────────────────────────────────────────
合计 ≈ 18BTDF 3DF
注意 gated MLP 是 3DF 不是 2DF。老式 Transformer（GPT-2 那种）是 2DF。现在主流（LLaMA / DeepSeek / Qwen）都用 gated，所以是 3DF。
3.2 Attention 部分（QKVO 投影）
操作 训练 FLOPs 参数
──────────────────────────────────────────────────────────
A[B,T,D] · W_Q[D,N,H] 6BTDNH DNH
A[B,T,D] · W_K[D,K,H] 6BTDKH DKH
A[B,T,D] · W_V[D,K,H] 6BTDKH DKH
A[B,T,N,H] · W_O[N,H,D] 6BTDNH DNH
──────────────────────────────────────────────────────────
合计 12BTD(N+K)H 2D(N+K)H
GQA 的好处直接体现在这里：当 K << N 时（Llama 3 70B 是 N=64, K=8）：
参数: 2D(N+K)H = 2D·D + 2D·(K/N)·D
MHA (K=N): 4D²
GQA (K=N/8): 2.25D² → 省 44%
3.3 Attention 的 dot-product 部分（关键！）
这部分常被初学者漏掉，也是长上下文的核心:
操作 训练 FLOPs
────────────────────────────────────────────────────────
Q[B,T,K,G,H] · K[B,S,K,H] → [B,T,S,K,G] 6BTSKGH = 6BTSNH
softmax_S ~O(BTSN)
S[B,T,S,K,G] · V[B,S,K,H] → [B,T,K,G,H] 6BTSKGH = 6BTSNH
────────────────────────────────────────────────────────
合计 ≈ 12BTSNH = 12BT²NH
(self-attention 时 S=T)
注意这里是 T² 而不是 T——这是 attention 唯一的二次项。
causal mask 减半：实际部署用 causal mask 时，有效计算量只有一半（下三角部分）。但要用 attention kernel 才能实际省下来，naive einsum 还是算全量。
3.4 汇总表
┌─────────────────────────────────────────────────────────┐
│ Component │ Params/layer │ Train FLOPs/layer │
├─────────────────┼──────────────┼────────────────────────┤
│ MLP (gated) │ 3DF │ 18BTDF │
│ Attention QKVO │ 2D(N+K)H │ 12BTD(N+K)H │
│ Attention dot │ 0 │ 12BT²NH │
│ LayerNorm ×2 │ 2D │ ~O(BTD) │
│ Unembed (全局) │ DV │ 6BTDV │
└─────────────────────────────────────────────────────────┘
参数量主导者是 MLP。用 MHA（K=N）算比例:
MLP : Attention = 3DF : 4D²
当 F = 4D 时: = 12D² : 4D² = 3 : 1
→ attention 只占参数量的 1/4!
四、Attention 什么时候开始贵过 FFN
这是原书里最漂亮的一个结论，面试价值极高。
问题：我们说 Transformer 是 matmul 主导的，但 attention 有 T² 项。什么时候 T² 会变成主要开销?
设定: F = 4D, N·H = D, N = K (MHA)
attention dot-product FLOPs = 12BT²NH = 12BT²D
所有 matmul FLOPs (含投影) = 18BTDF + 24BTDNH
= 18BTD(4D) + 24BTD·D
= 72BTD² + 24BTD²
= 96BTD²
比值:
attn / matmul = 12BT²D / 96BTD² = T / (8D)
黄金结论:
┌──────────────────────────────────────────────┐
│ attention FLOPs 开始超过所有 matmul FLOPs │
│ 当且仅当 T > 8D │
└──────────────────────────────────────────────┘
实际数字（Llama 3 70B, D=8192）:
T_critical = 8 × 8192 = 65,536 tokens
T = 4,096 → attn/matmul = 0.06 (attention 几乎免费)
T = 32,768 → attn/matmul = 0.50 (开始有感)
T = 65,536 → attn/matmul = 1.00 (打平!)
T = 131,072 → attn/matmul = 2.00 (attention 翻倍贵)
T = 1,000,000 → attn/matmul = 15.3 (attention 完全主导)
这解释了一个反直觉现象:大模型的长上下文其实比小模型容易。
Gemma-27B (D=4608): T_critical = 37k
Llama-70B (D=8192): T_critical = 65k
Llama-405B (D=16384): T_critical = 131k
D 越大,attention 的二次项越晚成为瓶颈
因为 FFN 是 O(D²),attention 是 O(T²),D 大相当于"抬高门槛"
注意：这个结论只针对 FLOPs。显存上 attention 的 [B,T,S,N] 矩阵仍然是 O(T²)，这是另一回事（FlashAttention 解决的正是这个）。
五、完整核算:Llama 3 70B
把公式套到真实模型。这是面试白板题的标准演练。
5.1 参数
L=80, D=8192, F=28672, V=128256, N=64, K=8, H=128
MLP/layer = 3 × 8192 × 28672 = 704.6 M
Attn/layer = D·N·H + 2·D·K·H + N·H·D
= 8192·8192 + 2·8192·1024 + 8192·8192
= 67.1M + 16.8M + 67.1M = 151.0 M
LN/layer = 2 × 8192 = 16 K (忽略)
每层合计 = 855.6 M
80 层 = 68.45 B
Embed+Unembed = 2 × 8192 × 128256 = 2.10 B
─────────────────────────────────
总计 = 70.55 B ✓ (与官方标称 70B 一致)
构成分析:
MLP 56.4 B (80%)
Attention 12.1 B (17%)
Embedding 2.1 B (3%)
如果不用 GQA 而用 MHA(K=64):
Attn = 4·D²·L = 21.5 B → 总参数 79.9 B
GQA 省了 9.4 B 参数!
5.2 KV Cache
KV Cache/token = 2 (K和V) × L × K × H × bytes
Llama 3 70B:
bf16: 2 × 80 × 8 × 128 × 2 = 320 KiB/token
fp8 : 2 × 80 × 8 × 128 × 1 = 160 KiB/token
单条序列:
8K 上下文 → 320K × 8192 = 2.68 GB
32K 上下文 → 320K × 32768 = 10.7 GB
128K 上下文→ 320K × 131072= 42.9 GB ← 单条就爆单卡!
这就是 GQA 存在的真实原因。对比 MHA:
MHA (K=64): KV/token = 2×80×64×128×2 = 2.5 MiB/token
8K 上下文 = 21 GB (单条序列就吃掉 1/4 张 H100)
128K = 336 GB (完全不可能)
GQA (K=8): 320 KiB/token — 省了 8 倍!
5.3 部署能撑多少并发
假设 8 卡 H100 跑 TP=8,bf16 权重:
显存预算:
每卡 80 GB
权重: 140 GB / 8 = 17.5 GB/卡
激活 + 框架开销: ~5 GB/卡
─────────────────────────────────
可用 KV Cache: ~55 GB/卡 × 8 = 440 GB
最大并发:
4K 上下文 → 440 / (0.32 × 4096) = 41 条
8K 上下文 → 440 / (0.32 × 8192) = 20 条
32K 上下文 → 440 / (0.32 × 32768) = 5 条
这张表就是你回答"能撑多少并发"的答案。想让并发翻倍?换 fp8 KV Cache（直接 x2）。
5.4 Decode 延迟的物理下界
Decode 每生成一个 token,要把全部权重从 HBM 读一遍:
140 GB / 3.35 TB/s = 41.8 ms (单卡放不下,纯理论)
TP=8 时:
每卡 17.5 GB / 3.35 TB/s = 5.2 ms
→ TPOT 理论下界 ~5.2 ms/token ≈ 190 tokens/s
→ 实际考虑 KV cache 读取和通信,约 15-25 ms (40-65 tokens/s)
注意:这个下界与 batch size 无关——因为 decode 时 batch 里的所有序列共享同一份权重读取。这就是为什么大 batch 时吞吐暴涨（一条权重读取服务 N 个请求）。
六、MoE 的账怎么算
MoE 的记账方式完全不同，是面试的重点。
6.1 参数 vs 激活参数
设定: E 个 expert,k 个被激活
总参数 = E × (dense MLP 参数) × L ← 显存看这个
激活参数 = k × (dense MLP 参数) × L ← FLOPs 看这个
稀疏度 = E / k (通常 8-64)
DeepSeek V3 实例:
总参数: 671 B
激活参数: 37 B
E = 256 (routed) + 1 (shared)
k = 8
FLOPs 按 37B 算,不是 671B!
显存按 671B 算,不是 37B!
─────────────────────────────────────────
这就是 MoE 的核心 trade-off:
用 18 倍显存换 2 倍算力效率
6.2 MoE 的 roofline:激活很稀疏时的致命问题
MoE 有个 roofline 陷阱——你以为省了 FLOPs,但实际上更容易 memory-bound。
Dense MLP:
FLOPs = 2·B·D·F Bytes = D·F
Intensity = 2B
MoE MLP (E experts, top-k):
FLOPs = 2·k·B·D·F (只算激活的)
Bytes = E·D·F (但要把所有 expert 权重都读出来!)
─────────────────────
Intensity = 2kB/E ← 除以了 E!
临界 batch size 变化:
compute-bound 条件: 2kB/E > 240 (TPU v5e)
⇔ B > 120 × E / k
DeepSeek V3 (E=256, k=8):
B > 120 × 256 / 8 = 3,840 tokens
→ 推理时 batch 要到 3840 tokens 才 compute-bound!
→ 这就是为什么 MoE 推理极难优化
对比 dense 模型只要 240，MoE 要 3840——差了 16 倍。这就是 MoE 推理的困难本质。
6.3 MoE 的通信账
每层 MoE 需要 2 次 AllToAll:
1. dispatch: 把 token 送到持有目标 expert 的卡
2. combine: 把结果送回来
通信量 (每层):
每次 AllToAll ≈ V / (4W) (环形拓扑,比 AllGather 便宜 4 倍)
L 层总通信 = 2L × V/(4W)
对比 TP 的 AllReduce:
4L × 2V/W (每层4次,每次2V/W)
─────────────────────────────
TP 通信是 MoE AllToAll 的 16 倍!
但 TP 走 NVLink (900 GB/s), MoE 走 IB (50 GB/s)
实际耗时反而 MoE 更差
七、Gradient Checkpointing:显存的救命稻草
这一节讲训练时激活值显存——它常常是 OOM 的真正原因。
7.1 激活值有多恐怖
B·T = 4M tokens, L = 64, D = 8192
完整保存所有中间激活:
每层约 20 个中间张量 (Transformer 图里每个节点)
sizeof = 2 × 20 × B·T × D × L
= 2 × 20 × 4e6 × 8192 × 64
= 84 TB !!
84 TB 是什么概念：一个 64 卡节点（8×H100×8 机）总共才 5 TB HBM。完全不可能。
7.2 三种 checkpointing 策略
┌───────────────────────────────────────────────────────────┐
│ 策略 │ 保存什么 │ 额外 FLOPs │ 显存 │
├────────────────────┼─────────────────┼─────────────┼───────┤
│ 无 checkpointing │ 全部激活 (20/层)│ 0% │ 84 TB │
│ Big matmuls only │ 7 个 matmul 输出│ ~10% │ 29 TB │
│ Block remat (全量) │ 只存层输入(1/层)│ ~33% │ 4.2TB │
└───────────────────────────────────────────────────────────┘
Block remat 的代价：
原来: 6 × params × tokens
block remat: 8 × params × tokens
↑
多一次完整 forward 重算
记忆口诀:
6x → 8x (block remat)
用 33% 的额外算力,换 20 倍的显存节省
7.3 各框架的做法
Megatron: --recompute-granularity full/selective
selective = 只重算 attention 部分(最省的选择)
DeepSpeed: activation_checkpointing
PyTorch: torch.utils.checkpoint.checkpoint()
JAX: jax.checkpoint / jax.remat
实战推荐：selective checkpointing（只重算 attention）通常是性价比最高的——attention 的中间张量最大（[B,T,S,N]）但重算成本相对低。
八、FlashAttention 为什么改变了账本
原书附录讲得很细，我提炼关键。
8.1 传统 attention 的问题不是 FLOPs 是显存
naive attention:
Step 1: S = Q · K^T → 物化 [B, T, S, N] 矩阵
T=128K, N=64, B=1 时: 128K × 128K × 64 × 2 = 2 TB !!
Step 2: P = softmax(S) → 又一个 [B,T,S,N]
Step 3: O = P · V
两个 [B,T,S,N] 写到 HBM 再读回来 → 显存爆炸 + 带宽爆炸
8.2 online softmax 的核心恒等式
标准 softmax:
Attn(Q,K,V) = Σᵢ [exp(Q·Kᵢ - max) · Vᵢ] / Σₗ exp(Q·Kₗ - max)
online softmax 的关键恒等式:
Σᵢ exp(aᵢ + b) = exp(b) · Σᵢ exp(aᵢ)
─────────────────────────────────────
意味着: 可以分块算,最后用 max 差值 rescale 累加!
分块合并公式（两块 K/V 为例）:
L¹ = Σᵢ exp(Q·Kᵢ¹ - M¹) ← 第一块的 local sum
L² = Σᵢ exp(Q·Kᵢ² - M²) ← 第二块的 local sum
M_combined = max(M¹, M²)
L_combined = exp(M¹ - M_combined) · L¹
+ exp(M² - M_combined) · L²
↑ 把第一块的 sum 按新 max rescale
这样只需要维护 3 个 running 状态:
M: running max ← 每个 Q 位置一个标量
L: running sum-of-exp ← 每个 Q 位置一个标量
O: running output ← 每个 Q 位置一个 D 维向量
显存: O(B·T·H) 而不是 O(B·T·S·H) → 从 quadratic 到 linear!
8.3 反向传播的巧妙恒等式
FlashAttention 反向时有一个关键技巧，值得单独讲（这是 ring attention 等方案的基础）:
定义 Sᵢⱼ = exp(τ qᵢ·kⱼ) / Σₗ exp(τ qᵢ·kₗ) (softmax 权重)
需要算: dSᵢⱼ (从 dO 和 V 反推)
dSᵢⱼ = Σd dOᵢd · Vⱼd
为了传给 Q 和 K 需要:
d(qᵢ·kⱼ) = (dSᵢⱼ - Sᵢⱼ · Σⱼ dSᵢⱼ) · Sᵢⱼ
↑
这一项原本需要对完整 S 长度求和 (贵!)
恒等式巧妙之处:
Σⱼ Sᵢⱼ · dSᵢⱼ = Σⱼ [exp(τqᵢ·kⱼ)/Σₖexp(...)] · Σd dOᵢd·Vⱼd
= Σd dOᵢd · Σⱼ Sᵢⱼ Vⱼd
= Σd dOᵢd · Oᵢd ← O 已经在 forward 算过了!
= dOᵢ · Oᵢ
→ 原来需要 O(S) 的求和,变成 O(D) 的局部点积!
这个恒等式是 ring attention / context parallel 能在反向传播时正确分块的基础。
九、面试锚点
Q1: 6 × params × tokens 这个公式怎么推导?
A: 单个矩阵乘 C = A·B，fwd 是 2NPM FLOPs。bwd 需要两个梯度:dL/dB = A^T·dL/dC（2NPM）和 dL/dA = dL/dC·B^T（2NPM）。总 6NPM，而 PM 就是参数量，N 是 token 数。所以 6 × params × tokens。推论:训练是推理的 3 倍（fwd 2NPM vs 训练 6NPM）。
Q2: 怎么快速算 Transformer 参数量?
A: 主项三条: MLP = 3DF·L（gated 是 3，老式是 2）; Attention(QKVO) = 2D(N+K)H·L（MHA 时是 4D²L）; Embedding+Unembed = 2DV（通常 tied，只算一份）。MLP 占约 75-80%，attention 占 ~17%，embedding 占 ~3%。
Q3: attention 什么时候开始比 FFN 更贵?
A: FLOPs 上，dot-product attention 与所有 matmul 的比值 = T/(8D)（假设 F=4D, N=K）。所以 T > 8D 时 attention 主导。Llama 70B（D=8192）→ 65K tokens;Gemma-27B（D=4608）→ 37K。反直觉结论:大模型的长上下文比小模型容易，因为 D 大把门槛抬高了。
Q4: KV Cache 怎么算?Llama 70B 一条 8K 序列占多少?
A: 2 × L × K × H × bytes。Llama 70B（L=80, K=8, H=128）:bf16 = 320 KiB/token，8K 序列 = 2.68 GB。128K = 42.9 GB（超单卡）。这就是 GQA 的动机——MHA（K=64）会是 2.5 MiB/token，8 倍差距。
Q5: 8 卡 H100 跑 Llama 70B bf16，能撑多少并发?
A: 每卡 80GB，权重 140/8=17.5GB，激活+开销约 5GB，剩 ~55GB/卡 × 8 = 440GB 给 KV Cache。4K 上下文 → 41 条；8K → 20 条；32K → 5 条。想翻倍就用 fp8 KV Cache。
Q6: MoE 的参数和 FLOPs 怎么算?
A: 显存按总参数算（所有 expert 都在显存里），FLOPs 按激活参数算（每个 token 只走 k 个 expert）。DeepSeek V3:总 671B，激活 37B。训练 FLOPs = 6 × 37e9 × tokens，但显存要按下 671B 规划。
Q7: MoE 推理为什么这么难优化?
A: roofline 上，MoE 的 intensity = 2kB/E，被 expert 数 E 除了一次。临界 batch = 120×E/k。DeepSeek V3（E=256, k=8）需要 3840 tokens/GPU 才 compute-bound，而 dense 模型只要 240——差 16 倍。加上 AllToAll 通信和负载不均衡，MoE 推理优化极难。
Q8: Gradient checkpointing 的代价是什么?
A: 用算力换显存。完整保存所有激活要 84TB（4M tokens, 64 层, D=8192），完全不可能。Block remat（只存每层输入）把显存降到 4.2TB，但 FLOPs 从 6×params×tokens 涨到 8×params×tokens（+33%）。Selective checkpointing（只重算 attention）通常性价比最高——attention 中间张量最大但重算便宜。
Q9（进阶）: FlashAttention 到底省了什么?
A: 省的不是 FLOPs 是 HBM 访问。naive attention 要物化 [B,T,S,N] 矩阵并写回 HBM（T=128K 时 2TB）。FlashAttention 用 online softmax 恒等式（Σexp(a+b)=exp(b)Σexp(a)）分块计算，只需维护 3 个 running 状态（M/L/O），显存从 O(T²) 降到 O(T)。FLOPs 完全相同，但 intensity 从远低于临界提升到接近临界——这就是它快的真正原因。
Q10（进阶）: 为什么 MoE 的 top-k 路由用 AllToAll 而不是 AllReduce?
A: 语义不同。AllReduce 是"每卡算部分和然后求和"——前提是所有卡在算同一个结果。MoE 需要的是"重新分组"：把 batch 里属于 expert_i 的 token 送到持有 expert_i 的卡。这是"搬下标"操作，AllToAll 恰好对应（[B_X, E] → [B, E_X]）。强行 AllReduce 需要大量 padding/masking，通信量和显存都爆。
十、一张速查卡
┌──────────────────────────────────────────────────────────┐
│ Transformer 账本速查 (gated MLP, GQA) │
├──────────────────────────────────────────────────────────┤
│ 参数量 │
│ MLP : 3DF per layer │
│ Attention : 2D(N+K)H per layer (MHA: 4D²) │
│ LayerNorm : 2D per layer (忽略) │
│ Embedding : 2DV total (通常 weight tying, 算一份) │
│ │
│ FLOPs (训练) │
│ MLP : 18BTDF per layer │
│ Attn proj : 12BTD(N+K)H per layer │
│ Attn dot : 12BT²NH per layer ← 唯一的二次项 │
│ 总计 : ≈ 6 × params × tokens (T < 8D 时成立) │
│ │
│ 显存 │
│ 权重 bf16 : 2 bytes/param │
│ 梯度 bf16 : 2 bytes/param │
│ Adam 状态 : 12 bytes/param (m + v + master, fp32) │
│ → 训练总计 : 16 bytes/param (纯 DP) │
│ KV Cache : 2 × L × K × H × bytes per token │
│ 激活值 : ~20 × BTDL × 2 bytes (要 checkpoint!) │
│ │
│ 关键临界值 │
│ attention 主导 : T > 8D │
│ matmul compute-bound : batch > 240 tokens │
│ MoE compute-bound : batch > 120·E/k tokens │
└──────────────────────────────────────────────────────────┘
一句话记住这一篇
Transformer 的账本三条主线:参数上 MLP 占 80%（3DF/层，attention 只有 2D(N+K)H）;FLOPs 上 6×params×tokens 是 fwd 2x + bwd 4x 的严格结果，但 T>8D 后 attention 的二次项反超;显存上训练要 16 bytes/param，推理真正的大头是 KV Cache（2LKH per token）。MoE 的账反着记——显存看总参数，FLOPs 看激活参数，但临界 batch 要乘 E/k，这就是 MoE 推理难优化的数学根源。
系列下一篇预告:《LLaMA 3 训练实操:几张卡、几天、几刀》——把前四篇的公式全用上,算一次真实的 70B 训练预算:选 TP/PP/DP 组合、估 step time、算 MFU、最后得出美元成本。
配套阅读:JAX Scaling Book Ch4 原书章节。相关文章:Roofline 分析、Sharded Matmul、分布式 Transformer 深度解析、MoE 推理完整机制。
Written by
Zoe
AI Infra Engineer · LLM Serving · GPU/RDMA · 造工具的偏执狂