Llama 3 8B 只有 80 亿参数，可用 BF16 混合精度加 Adam 训练它，光权重、梯度和优化器状态就要约 128 GB，一张 80 GB 的 H100 放不下，激活值还没算。数据并行、ZeRO、张量并行、流水线并行，本质上都在回答同一个问题：一步训练要花多少显存、多少计算、多少通信，怎么在这三者之间做交换。哪怕只有一张卡，也有三本账要算：
- 显存账：权重、梯度、优化器状态、激活值各占多少字节，怎么推出来的；
- 计算账：一次前向和反向到底做了多少次浮点运算（FLOPs），著名的 \(6ND\) 从哪来，注意力的二次项什么时候才重要；
- 效率账：MFU 和 HFU 怎么定义，怎么用它们估算训练要多少天。
1. 三个约束：显存、计算、通信
训练一个大模型，始终在和三个约束打交道：
- 显存是硬约束。一步训练放不进显存，就根本跑不起来（OOM，out of memory）。
- 计算效率：希望硬件绝大部分时间在做计算，而不是在搬数据、等别的卡。
- 通信开销：卡与卡之间传数据时，计算单元常常在空等。要尽量用快的链路（机内）而少用慢的链路（机间），并且让通信和计算重叠。
几乎每一种技术都是在这三者之间做交换：激活重计算用计算换显存；张量并行用通信换显存；梯度累积用时间换显存。单卡上只有前两个约束；通信要到多卡才出现，量化它的工具是 Roofline。
2. 一步训练，和以 token 计的 batch size
单卡上的一步训练（一个 step）有三个阶段：
- 前向（forward）：输入穿过模型，得到输出和 loss；
- 反向（backward）：从 loss 出发用链式法则算出每个参数的梯度；
- 优化器更新（optimizer step）：用梯度更新参数，例如 Adam。
Batch size（记作 \(bs\)）是一步里用多少条样本估计梯度。它同时影响收敛和吞吐：太小，梯度噪声大；太大，每个 token 被利用得不充分，同样的 token 数换来的 loss 下降变少。好在最终效果对 batch size 在最优值附近并不敏感，所以工程上有不少调整余地。有的训练还会在过程中逐步增大 batch：DeepSeek-V3 在前 469B token 里把 batch 从 3,072 条序列逐渐增大到 15,360 条，之后保持不变；Llama 3 405B 先用 4M token 的 batch（序列长 4,096），训练 252M token 后翻倍到 8M（序列长 8,192），训练到 2.87T token 时再翻倍到 16M。
预训练圈子里习惯用 token 数 报告 batch size，因为它和序列长度无关：
$$bst = bs \times seq$$这里 \(seq\) 是序列长度。近年 LLM 预训练的 batch 大致在每步 4M 到 60M token 之间：Llama 1 用约 4M token 的 batch 训练了 1.4T token，DeepSeek-V3 最终用约 63M token（15,360 条 × 4,096）的 batch 训练了 14.8T token。问题马上就来了：一张 80 GB 的 H100 显然装不下 4M token 的激活值。为了搞清楚怎么办，先看显存到底花在哪。
3. 显存里住着什么
训练时显存里主要有四样东西：
- 模型权重（weights / parameters）；
- 梯度（gradients），和权重一一对应，形状相同；
- 优化器状态（optimizer states），例如 Adam 的一阶矩 \(m\) 和二阶矩 \(v\)，也和权重形状相同；
- 激活值（activations）：前向时产生、反向时要用到的中间张量。
前三样只和模型大小有关，是「静态」的；激活值和输入有关，随 batch size 和序列长度变化。此外还有一些难以精确计算的杂项：CUDA 上下文和内核本身要占 1–2 GB（可以执行一句 torch.ones((1, 1)).to("cuda")，再用 nvidia-smi 看看占了多少）；中间缓冲区和显存碎片也会吃掉一些。下面的估算都忽略这些常数项。
一步之内显存怎么变化
显存占用并不是一个静态数字，在一步之内起伏很大。用 PyTorch profiler 看一步训练，形状通常是：前向时激活值快速堆高；反向开始后，梯度一层层生成，同时用过的激活值一层层释放；最后优化器用全部梯度更新参数和优化器状态，然后进入下一步。
还有一个容易踩的坑：第一步能跑，第二步才 OOM。原因是 Adam 的状态在第一次调用 optimizer.step() 时才分配（PyTorch 是惰性创建的），所以从第二步开始，优化器状态和激活值同时在显存里，峰值更高。另外，第一步的显存曲线形状也和后面不同：PyTorch 的缓存分配器在第一步里做了大量预分配，之后的步就不用再找空闲块了。下面的动画把前两步画了出来。
一个 4 层的模型，每层参数量相同，按 BF16 混合精度 + Adam 计（每个参数：权重 2 + FP32 主权重 4 = 6 份、梯度 2 份、优化器状态 8 份）；每层激活值的大小和虚线表示的显存上限都是示意值。
4. 数参数：从公式到 Llama 3
静态显存 = 参数量 × 每个参数的字节数。先数参数。记号：隐藏维度 \(h\)，层数 \(L\)，词表大小 \(V\)，注意力头数 \(a\)，每头维度 \(d_h = h/a\)，MLP 中间维度 \(F\)。
经典 GPT 结构
先看最经典的 GPT 结构：线性层带偏置，用 LayerNorm，MLP 中间维度是 \(4h\)，输出层和输入嵌入共享权重。参数量（不算位置嵌入）是：
$$N = h\,V + L\,(12h^2 + 13h) + 2h$$逐项拆开就是：
- 词嵌入 \(hV\)（输出层与之共享，不另算）；
- 注意力：\(W_Q, W_K, W_V, W_O\) 四个 \(h\times h\) 矩阵，\(4h^2\)，加上 4 个偏置 \(4h\)；
- MLP：\(h \to 4h \to h\) 两个矩阵 \(8h^2\)，偏置 \(4h + h = 5h\)；
- 两个 LayerNorm，每个有缩放和偏移 \(2h\)，共 \(4h\)；
- 每层合计 \(12h^2 + 13h\)，最后再加一个 LayerNorm 的 \(2h\)。
当 \(h\) 大时，\(h^2\) 项占绝对主导，所以常用的近似是 \(N \approx 12Lh^2\)。
Llama 结构：GQA + SwiGLU + RMSNorm
现代模型有三个变化：（1）注意力用 GQA（grouped-query attention），K、V 只有 \(n_{kv}\) 个头，于是 \(W_K, W_V\) 是 \(h \times n_{kv} d_h\)；（2）MLP 用带门控的 SwiGLU，有三个矩阵 \(W_\text{gate}, W_\text{up}\in\mathbb{R}^{h\times F}\)、\(W_\text{down}\in\mathbb{R}^{F\times h}\)；（3）RMSNorm 只有缩放参数，没有偏置，线性层也不带偏置。于是每层：
$$N_\text{layer} = \underbrace{2h\cdot a d_h + 2h\cdot n_{kv} d_h}_{\text{注意力}} + \underbrace{3hF}_{\text{MLP}} + \underbrace{2h}_{\text{RMSNorm}}$$整个模型再加上输入嵌入 \(Vh\)、不共享的输出层 \(Vh\)、最后一个 RMSNorm \(h\)。拿 Llama 3 的三个尺寸代进去（三者的词表都是 128,256，都是 \(n_{kv}=8\)、\(d_h = 128\)，输出层都不与嵌入共享）：
| 8B | 70B | 405B | |
|---|---|---|---|
| 层数 \(L\) | 32 | 80 | 126 |
| 隐藏维度 \(h\) | 4,096 | 8,192 | 16,384 |
| MLP 维度 \(F\) | 14,336 | 28,672 | 53,248 |
| Q 头数 \(a\) / KV 头数 | 32 / 8 | 64 / 8 | 128 / 8 |
| 每层注意力参数 | 41.9M | 151.0M | 570.4M |
| 每层 MLP 参数 | 176.2M | 704.6M | 2,617.2M |
| 嵌入 + 输出层 | 1.05B | 2.10B | 4.20B |
| 总参数 \(N\) | 8.03B | 70.55B | 405.85B |
几个观察：MLP 占每层参数的约 80%（GQA 让注意力更便宜了）；8B 模型里光嵌入和输出层就有 1.05B，占 13%，这是大词表在小模型上的代价；到 405B，它们只占 1% 左右。
5. 每个参数几个字节：16 字节规则及其变体
精度决定每个数占几个字节：FP32 是 4 字节，BF16/FP16 是 2 字节，FP8 是 1 字节。下面以最常用的 Adam（或 AdamW）为例，它为每个参数存两个状态：动量 \(m\) 和方差 \(v\)。
全 FP32 训练
$$m_\text{params} = 4N,\quad m_\text{grad} = 4N,\quad m_\text{opt} = (4+4)N \;\Rightarrow\; 16N$$BF16 混合精度 + FP32 主权重
现在默认的做法是「混合精度」：前向和反向的矩阵乘法用 BF16 算，权重和梯度也以 BF16 存一份；但为了数值稳定，另存一份 FP32 的权重拷贝，叫主权重（master weights），优化器在这份 FP32 拷贝上更新，然后再转成 BF16 给下一步用。Adam 的 \(m, v\) 也用 FP32 存：
$$m_\text{params} = 2N,\quad m_\text{grad} = 2N,\quad m_\text{params,fp32} = 4N,\quad m_\text{opt} = (4+4)N \;\Rightarrow\; 16N$$有些框架（例如 Nanotron）还会额外保留一份 FP32 的梯度用来累加，因为 BF16 只有 7 位尾数，把很多个小梯度加在一起时会丢精度。这样再加 \(4N\)，总共 20 字节/参数。
一个反直觉的事实：混合精度本身并不省权重相关的显存，只是把 16 字节重新分配了一下，带 FP32 梯度累加时甚至更多。它的好处在别处：（1）BF16 矩阵乘法在 GPU 上快得多（H100 上 BF16 Tensor Core 的峰值是 TF32 的 2 倍、FP32 CUDA core 的约 15 倍）；（2）激活值用 BF16 存，激活显存减半——而激活值往往是大头。
| 配置 | 权重 | 梯度 | 主权重 | 优化器 | 合计 / 参数 |
|---|---|---|---|---|---|
| FP32 + Adam | 4 | 4 | – | 8 | 16 |
| BF16 混合 + FP32 主权重 + Adam | 2 | 2 | 4 | 8 | 16 |
| 同上，梯度以 FP32 累加 | 2 | 2 + 4 | 4 | 8 | 20 |
| FP32 + 带动量的 SGD | 4 | 4 | – | 4 | 12 |
| 只做推理（BF16） | 2 | – | – | – | 2 |
后两行是为了对照：优化器状态的大小取决于优化器本身（带动量的 SGD 只存一个状态），而推理只需要权重。用 FP8 训练能进一步压低，但稳定性更难保证（见混合精度与 FP8）。
按 16 和 20 字节/参数，几个常见规模的静态显存是：
| 参数量 | 16 字节/参数 | 20 字节/参数 |
|---|---|---|
| 1B | 16 GB | 20 GB |
| 7B | 112 GB | 140 GB |
| Llama 3 8B（8.03B） | 128 GB | 161 GB |
| 70B | 1,120 GB | 1,400 GB |
| 405B | 6,480 GB | 8,100 GB |
结论很直接：7B 往上，光权重、梯度和优化器状态就已经超过单张 H100 的 80 GB。还没算激活值。这就是为什么后面需要 ZeRO/FSDP 把这些状态切到多张卡上。
6. 激活显存：逐项推导
为什么反向需要存激活值
考虑一个线性层 \(Y = XW\)，其中 \(X\) 是 \([s, h_\text{in}]\)，\(W\) 是 \([h_\text{in}, h_\text{out}]\)。反向时拿到上游梯度 \(\partial \mathcal{L}/\partial Y\)，需要算两个东西：
$$\frac{\partial \mathcal{L}}{\partial W} = X^\top \frac{\partial \mathcal{L}}{\partial Y}, \qquad \frac{\partial \mathcal{L}}{\partial X} = \frac{\partial \mathcal{L}}{\partial Y} W^\top$$第一个式子需要前向时的输入 \(X\)。所以每个矩阵乘法都要把自己的输入留到反向。逐元素的非线性函数也一样：\(y = \sigma(x)\) 的反向 \(\partial\mathcal{L}/\partial x = \sigma'(x)\cdot \partial\mathcal{L}/\partial y\) 需要 \(x\)（或者 \(y\)，看函数）。这些在前向时产生、留到反向才用的张量，就是激活值。
GPT 结构的账：\(sbh(34 + 5as/h)\)
知道了「每个矩阵乘法要留输入」这条规则，就可以对一层逐项记账。先算经典的 GPT 结构（带 dropout、朴素注意力），这也是 Korthikanti 等人（2022）给出的经典结果。设序列长 \(s\)、micro-batch 大小 \(b\)（一次前向里的样本数），激活值用 2 字节存，dropout 的掩码用 1 字节存：
| 位置 | 要存的张量 | 字节 |
|---|---|---|
| 注意力 | Q/K/V 投影的共同输入 | \(2sbh\) |
| \(QK^\top\) 需要 Q 和 K | \(4sbh\) | |
| softmax 的输出 | \(2as^2b\) | |
| softmax 后 dropout 的掩码 | \(as^2b\) | |
| 乘 V：dropout 后的概率和 V | \(2as^2b + 2sbh\) | |
| 输出投影的输入 + 其后 dropout 掩码 | \(2sbh + sbh\) | |
| MLP（\(F=4h\)） | 第一个线性层的输入 | \(2sbh\) |
| GeLU 的输入、第二个线性层的输入 | \(8sbh + 8sbh\) | |
| dropout 掩码 | \(sbh\) | |
| LayerNorm | 两个 LayerNorm 各存输入 | \(4sbh\) |
注意力块合计 \(11sbh + 5as^2b\)，MLP \(19sbh\)，LayerNorm \(4sbh\)，加起来：
$$m_\text{act/layer} = sbh\left(34 + \frac{5as}{h}\right)\ \text{字节},\qquad m_\text{act} = L\cdot sbh\left(34 + \frac{5as}{h}\right)$$两点要看清楚：（1）激活显存和 \(b\) 成线性、和 \(s\) 有二次项，来源是 \(s\times s\) 的注意力矩阵，每个头一份；（2）和权重不同，它跟着输入涨。把具体模型代进去就会发现：\(b=1\) 时，短序列的激活值几乎可以忽略，从 2k–4k token 开始变得显著，长序列时远远超过权重相关的部分（第 7 节计算器里的曲线图可以直接看到）。
换成现代模型：没有 dropout，有 FlashAttention、GQA 和 SwiGLU
上面是 2022 年 GPT-3 风格模型的账。今天的 Llama 类模型有四处不同：
- 没有 dropout：预训练大模型基本不用 dropout，掩码项全部消失。
- FlashAttention：不把 \(s\times s\) 的注意力矩阵写回显存，反向时分块重算。它只需要存 Q、K、V、输出 O，以及每行 softmax 的 log-sum-exp（FP32，每个头每个位置一个数，\(4abs\) 字节，可忽略）。整个 \(as^2\) 项消失了。
- GQA：K、V 只有 \(n_{kv}\) 个头，记 \(h_{kv} = n_{kv} d_h\)，K、V 共 \(4sbh_{kv}\) 字节。
- SwiGLU：\(\text{down}(\text{SiLU}(xW_\text{gate}) \odot xW_\text{up})\)。在不融合内核的 PyTorch 写法里，要存 gate 的输出（SiLU 的输入）、up 的输出、SiLU 的输出（逐元素乘的另一个输入）、乘积（down 的输入），共四个 \([s,b,F]\) 张量。
加起来，每层（BF16，激活每个数 2 字节）：
$$m_\text{act/layer} \approx 2sb\,\big(\underbrace{2h}_{\text{两个 norm 输入}} + \underbrace{h}_{\text{QKV 输入}} + \underbrace{h + 2h_{kv}}_{Q,K,V} + \underbrace{h}_{O} + \underbrace{h}_{\text{MLP 输入}} + \underbrace{4F}_{\text{SwiGLU}}\big) = 2sb\,(6h + 2h_{kv} + 4F)$$验算一下：把它套回 GPT 结构（\(h_{kv}=h\)，GeLU 只有输入和输出两个 \(4h\) 张量，即把 \(4F\) 换成 \(2\cdot 4h\)），得到 \(2sb\cdot 16h = 32sbh\)，正好是前面的 34 减去两个 dropout 掩码（\(sbh\) 各一个）。如果用朴素注意力（不用 FlashAttention），还要加上 softmax 的输出 \(2as^2b\)。这本账按不融合内核的 PyTorch 写法来记；把 SwiGLU 写成融合内核的框架，还能再少存一两个 \([s,b,F]\) 张量。
对 Llama 3 8B（\(h=4096, h_{kv}=1024, F=14336\)），括号里是 83,968，每个 token 每层约 168 KB。序列长 8,192、\(b=1\) 时，32 层合计约 44 GB；如果用朴素注意力，光 softmax 矩阵就有 \(32\times 2 \times 32 \times 8192^2\) 字节 ≈ 137 GB。FlashAttention 的意义一目了然。
还有一个容易忽略的大张量：输出层的 logits。它的形状是 \([s, b, V]\)，Llama 3 的 \(V = 128{,}256\)，如果按 FP32 算交叉熵，一条 8k 序列就要 \(4 \times 8192 \times 128256\) ≈ 4.2 GB。这也是为什么很多框架会把交叉熵分块算或者和输出层融合。
7. 交互：训练显存计算器
选模型（或改配置）、序列长度、micro-batch、精度和重计算方式，看单卡显存的构成。计算器假设整个模型放在一张卡上，没有任何并行切分。
改配置
激活值按上一节的 Llama 式推导：\(2sb(6h+2h_{kv}+4F)\) 字节/层（BF16；FP32 模式下按 4 字节），朴素注意力另加 softmax 输出 \(2as^2b\)；完全重算时每层只存输入 \(2sbh\)，再加反向重算一层时的工作集。logits 按 FP32 的 \(4sbV\) 计；「杂项」固定记 1.5 GB，代表 CUDA 上下文等。显存碎片和临时缓冲区没有计入：真实框架里，内核融合会让激活更少，碎片和缓冲区会让总占用更多。
拖一拖就能看到几件事：
- Llama 3 8B 在 BF16 混合精度下，静态部分就要约 128 GB，单张 H100 怎么调都放不下：7B 以上的模型必须把这些状态切到多张卡上。
- 把序列长度从 4k 拉到 32k，朴素注意力下激活值会爆炸（二次项），FlashAttention 下是线性增长。
- micro-batch 加倍，激活值严格加倍；权重相关的部分不动。
- 完全重算把激活值压到只剩每层一个 \([s,b,h]\) 张量，代价是多算一遍前向（下一节）。
8. 激活重计算：拿算力换显存
激活重计算（activation recomputation），也叫梯度检查点（gradient checkpointing）或再物化（rematerialization），思路是：前向时只在少数关键位置保存激活值（叫检查点），其余的丢掉；反向需要时，从最近的检查点重新跑一段前向把它们算回来。本质上是多做一部分前向计算，换更少的显存。
完全重算（full）
在每个 Transformer 层的边界设检查点，即只存每层的输入。反向到某层时，先用这个输入把该层前向重跑一遍，再做反向。等于在反向里多做一次完整的前向。按每个参数、每个 token 计，前向 2 FLOPs、重算 2、反向 4，训练 FLOPs 从 \(6ND\) 变成约 \(8ND\)（第 10 节会推导这个 2 : 4），多出三分之一；实际一步的耗时通常也会增加 30–40%。显存方面，激活值从每层几十个 \(sbh\) 降到每层一个 \(2sbh\)，再加上反向时一层的工作集。
选择性重算（selective）
完全重算有点浪费：并不是所有激活都同样「贵」。关键的观察是：注意力里 \(QK^\top\)、softmax、乘 V 这几步，激活值很大（带 \(as^2\) 项）但计算量相对小（每个元素只做很少的运算）；而线性层的激活值小、重算却贵。所以只丢掉注意力核心部分，保留其余。公式里的 \(5as/h\) 项正是来自这部分：对 GPT-3（\(a=96, s=2048, h=12288\)）\(5as/h = 80\)，比 34 还大。Korthikanti 等人在 GPT-3 上测得，这样能省 70% 的激活显存，只多 2.7% 的 FLOPs。选择性重算后每层的激活值变成 \(34\,sbh\)，和序列长度呈线性。
如果把一次迭代的 FLOPs 精确写出来（包括注意力和输出层），模型本身的 FLOPs 和选择性重算后硬件实际做的 FLOPs 分别是：
$$\text{model FLOPs} = 72BLsh^2\left(1 + \frac{s}{6h} + \frac{V}{12hL}\right),\qquad \text{hardware FLOPs} = 72BLsh^2\left(1 + \frac{s}{3h} + \frac{V}{12hL}\right)$$（这里 \(B\) 是 batch 大小。）多出来的就是注意力的两个 \(s^2\) 矩阵乘法在前向里再做一遍，二者之比约为 \(1 + s/(6h)\)。
FlashAttention 已经是选择性重算
FlashAttention 在反向里本来就重算注意力分数和概率矩阵，而不是存下来。所以用 FlashAttention 的人，其实已经在做选择性重算了。另外，重算虽然增加了 FLOPs，却减少了对显存的读写；在显存带宽紧张的 GPU 上，总体常常反而更快。DeepSeek-V3 走得更远：反向时重算所有的 RMSNorm 和 MLA（多头潜在注意力）的上投影，这些输出一律不存。
其他策略
也可以从计算图的角度看这件事。如果什么都不重算，一层要存的中间节点大约有 20 个 \([B,T,D]\) 大小的张量（例如 \(f(x) = \exp(g(x))\) 的反向既要 \(g(x)\) 又要 \(\exp(g(x))\)）。举个例子：\(B\cdot T = 4\text{M}\) token、\(L=64\)、\(D=8192\)，BF16 下要存约 \(2\times 20\times BTDL\) ≈ 84 TB 的激活值。两种常用策略：
- Block remat：每层只存输入（即上面的完全重算），上例降到约 4.2 TB，FLOPs 从 \(6ND\) 变成约 \(8ND\)；
- 只存大矩阵乘法的输出：反向时不用重算任何大矩阵乘法，但要重算激活函数和注意力的一部分，每层从约 20 个降到约 7 个。
在 JAX 里这些策略由 jax.checkpoint（jax.remat）的 policy 控制；PyTorch 里对应 torch.utils.checkpoint，Megatron 有 --recompute-granularity full/selective。
补充：为什么是 \(\sqrt{L}\) 个检查点
这是 Chen 等人 2016 年的经典分析（《Training Deep Nets with Sublinear Memory Cost》）。假设 \(L\) 层，每层激活大小相同记为 1。每隔 \(k\) 层存一个检查点，就要存 \(L/k\) 个检查点；反向重算一段时，这一段的 \(k\) 层激活同时在显存里。峰值约为 \(L/k + k\)，对 \(k\) 求最小得 \(k = \sqrt{L}\)，峰值 \(2\sqrt{L}\)，代价同样是多一次前向。Transformer 训练里更常见的是「每层一个检查点」，因为一层内部的激活远大于层输入，按层切天然合适。
重算还让硬件实际做的 FLOPs 变多了，所以衡量效率时要区分 HFU 和 MFU（第 13 节）。
9. 梯度累积
重计算压住了「每个样本」的激活值，但激活值仍和 batch size 成正比。梯度累积（gradient accumulation）解决这个问题：把一个大 batch 切成若干 micro-batch，依次对每个 micro-batch 做前向和反向，把梯度累加起来，最后才做一次优化器更新。
记每次前向的 batch 为 micro-batch size \(mbs\)，两次优化器更新之间处理的总 batch 为 global batch size \(gbs\)：
$$gbs = mbs \times grad\_acc$$（有数据并行时，还要再乘上数据并行的副本数。）
optimizer.zero_grad()
for i, micro in enumerate(split(batch, grad_acc)):
loss = model(micro).loss / grad_acc # 取平均，结果与 grad_acc 无关
loss.backward() # 梯度在 .grad 里累加
optimizer.step()要点：
- 显存：激活值只需容纳一个 micro-batch，与 \(gbs\) 无关。代价是要有一块在整步里一直存在的梯度缓冲区。不做累积时，梯度可以边算边被优化器消费、激活边算边释放，峰值反而可以更低。
- 时间：micro-batch 是串行执行的，所以累积本身不加速，只是让大 batch 放得下：用时间换显存。
- 结果等价：数学上，取平均后累积的梯度等于整个 batch 的梯度。一个小坑：如果 loss 是按「每个 micro-batch 内 token 的平均」算的，而不同 micro-batch 有效 token 数不同（比如有 padding 或按文档打包），简单地再除以 \(grad\_acc\) 就不等于全局按 token 平均。正确做法是累加 token 级 loss 之和，最后除以整个 batch 的有效 token 数。
- 它和重计算可以叠加使用。
最后还有一个关键观察：各个 micro-batch 的前向和反向是彼此独立的，只差输入不同——那为什么不放到多张卡上并行做？这就是数据并行：它就是并行版本的梯度累积。
10. 数 FLOPs：从 \(2mnk\) 开始
FLOP 指一次浮点运算（一次乘或一次加），FLOPs 是运算次数，FLOP/s 是每秒运算次数（硬件峰值用它衡量）。
点积、矩阵向量积、矩阵乘法
- 两个长 \(P\) 的向量做点积：\(P\) 次乘、约 \(P\) 次加，\(2P\) FLOPs。
- 矩阵 \(A[N,P]\) 乘向量 \(x[P]\)：\(N\) 个点积，\(2NP\)。
- \(A[N,P]\) 乘 \(B[P,M]\)：对 \(B\) 的每一列做一次矩阵向量积，\(2NPM\)。
推广到一般情形：两个高维数组做缩并（einsum）时，FLOPs = 2 × 所有出现维度的乘积，其中缩并维（两边都有、输出里没有，被求和掉的）和批维（两边都有、输出里也有）只算一次。例如 \(C[G,H,I,J,K,L]\cdot D[G,H,M,N,K,L]\to[G,H,I,J,M,N]\)，\(G,H\) 是批维，\(K,L\) 是缩并维，FLOPs 为 \(2GHIJMNKL\)。没有缩并维时（纯逐元素乘）就没有那个 2。
值得记住的一点：矩阵乘法的计算量是 \(O(N^3)\)，而数据量只有 \(O(N^2)\)。矩阵越大，越容易让计算单元「吃饱」。这正是深度学习架构由矩阵乘法主导的原因。
反向是前向的两倍
设 \(C = AB\)，\(A\) 是 \([N,P]\) 的输入激活，\(B\) 是 \([P,M]\) 的权重。反向需要两个梯度：
$$\frac{\partial \mathcal{L}}{\partial B} = A^\top \frac{\partial \mathcal{L}}{\partial C}\quad(\text{缩并 } N,\ 2NPM),\qquad \frac{\partial \mathcal{L}}{\partial A} = \frac{\partial \mathcal{L}}{\partial C} B^\top\quad(\text{缩并 } M,\ 2NPM)$$所以训练时一个矩阵乘法总共 \(6NPM\) FLOPs：前向 \(2NPM\)，反向 \(4NPM\)，前向 : 反向 = 1 : 2。\(PM\) 正好是这个权重矩阵的参数量，\(N\) 是 token 数，于是每个 token 对每个参数花 6 FLOPs。这就是著名的 \(6 \times \text{参数量} \times \text{token 数}\) 的最简形式。
注意第一层（嵌入层后的第一个线性层）严格说不需要算输入梯度 \(\partial\mathcal{L}/\partial A\)，因为没有更前面的可学习层需要它；但这只占一层的一部分，可以忽略。
11. Transformer 的总账：\(6ND\) 与注意力项
现在把整个 Transformer 的 FLOPs 加起来。为了让 einsum 的维度一目了然，换一套更紧凑的记号：batch \(B\)、序列长 \(T\)、模型维度 \(D\)（即前面的 \(h\)）、MLP 维度 \(F\)、Q 头数 \(N\)（即 \(a\)）、KV 头数 \(K\)、头维 \(H\)（即 \(d_h\)）、词表 \(V\)。下表都是训练 FLOPs（前向 + 反向，已乘 3），除以 3 就是只做前向的量。
MLP（门控）
| 运算 | 训练 FLOPs | 参数 |
|---|---|---|
| \(A[B,T,D]\cdot W_{in1}[D,F]\) | \(6BTDF\) | \(DF\) |
| \(A[B,T,D]\cdot W_{in2}[D,F]\) | \(6BTDF\) | \(DF\) |
| \(\sigma(A_{in1}) \odot A_{in2}\) | \(O(BTF)\) | |
| \(A[B,T,F]\cdot W_{out}[F,D]\) | \(6BTDF\) | \(DF\) |
| 合计 | \(\approx 18BTDF\) | \(3DF\) |
注意力投影（GQA）
| 运算 | 训练 FLOPs | 参数 |
|---|---|---|
| \(W_Q[D,N,H]\) | \(6BTDNH\) | \(DNH\) |
| \(W_K[D,K,H]\)、\(W_V[D,K,H]\) | \(2\times 6BTDKH\) | \(2DKH\) |
| \(W_O[N,H,D]\) | \(6BTDNH\) | \(DNH\) |
| 合计 | \(12BTD(N+K)H\) | \(2D(N+K)H\) |
注意力核心（点积注意力本身）
把 Q 重排成 \([B,T,K,G,H]\)（\(G = N/K\) 是每组的 Q 头数）。\(Q\cdot K^\top\) 在 \(B, K\) 上是批维、在 \(H\) 上缩并，得到 \([B,T,S,K,G]\)；softmax 是逐元素级的 \(O(BTSN)\)；再乘 V 在 \(S\) 上缩并。两次矩阵乘法各 \(6BTSKGH = 6BTSNH\)：
$$\text{attention FLOPs}_\text{train} \approx 12\,BTSNH = 12\,BT^2NH\quad(\text{自注意力 } S=T)$$它不含任何参数，所以不在「6 × 参数量」里。对因果掩码，有用的 FLOPs 只有一半（需要注意力内核真的跳过被掩掉的块才能拿到这一半）。
其他
每层两个 LayerNorm/RMSNorm 只有 \(O(BTD)\)，可以忽略。最后的输出层（unembedding）是一个 \(6BTDV\) 的矩阵乘法，参数 \(DV\)，同样遵守 6 倍规则。输入嵌入是查表，不是矩阵乘法，严格说不花 FLOPs。
汇总成 6ND
忽略注意力核心，全部层加起来：
$$\big(18BTDF + 12BTD(N+K)H\big)L = 6\cdot BT\cdot\big(3DF + 2D(N+K)H\big)L = 6 \times \text{token 数} \times \text{参数量}$$记 \(N_\text{params}\) 为参数量，\(D_\text{tokens}\) 为训练 token 数（这里的 \(D\) 和上面的模型维度不是一回事，这是领域里的历史记号），得到稠密 Transformer 训练算力的经验公式：
$$C_\text{train} \approx 6\,N_\text{params}\,D_\text{tokens}$$验证一下：Llama 3 405B 训练了 15.6T token，\(6 \times 405\text{B} \times 15.6\text{T} \approx 3.79\times 10^{25}\)，Llama 3 论文给的预训练算力正是 \(3.8\times 10^{25}\) FLOPs。
注意力项什么时候重要
假设典型配置 \(F = 4D\)、\(D = NH\)、\(N=K\)（MHA），注意力核心与所有矩阵乘法（含投影）之比是：
$$\frac{12BT^2NH}{18BTDF + 24BTDNH} = \frac{12BT^2D}{72BTD^2 + 24BTD^2} = \frac{T}{8D}$$所以只有 \(T > 8D\) 时注意力 FLOPs 才占主导。对 \(D \approx 8\text{k}\) 的模型，要到约 64k token；Gemma 2 27B（\(D = 4608\)）按这个估算大约在 37k token 时注意力开始占主导。这说明对大模型来说，注意力的二次方代价并不像直觉中那么可怕：模型越宽，MLP 越贵，注意力越不显眼。（另一个相关的结论：注意力核心 FLOPs 等于 QKVO 投影 FLOPs 的条件是 \(24BTDNH = 12BT^2NH\)，即 \(T = 2D\)。）
拿 Llama 3 8B 在 8k 序列上具体算：\(6N\) ≈ 48.2 GFLOPs/token；注意力核心按因果减半是 \(6\cdot T\cdot NH\cdot L = 6 \times 8192 \times 4096 \times 32\) ≈ 6.4 GFLOPs/token，大约是 \(6N\) 的 13%。405B 在同样长度下只有约 4%。
也有人把每 token 的训练 FLOPs 写成 \(6N + 12\,L\,T\,d_\text{attn}\)（\(d_\text{attn} = NH\)）的形式（不考虑因果减半），和上面是同一件事。
12. MoE 与 KV cache 的账
MoE
混合专家（Mixture of Experts）把每层的一个稠密 MLP 换成 \(E\) 个独立的 MLP（专家），每个 token 只路由到其中 \(k\) 个，通常 \(k \ll E\)；\(E/k\) 叫稀疏度，常见在 8 到 64 之间（例如 DeepSeek-V3 每层有 256 个路由专家，每个 token 选 8 个，另有 1 个所有 token 都经过的共享专家）。参数量按 \(E\) 增长，每个 token 实际经过的参数（激活参数）只按 \(k\) 增长。所以算 FLOPs 时要用激活参数：
$$C_\text{train}^\text{MoE} \approx 6\,N_\text{active}\,D_\text{tokens}$$而显存账要按总参数算——所有专家的权重、梯度和优化器状态都得存着。这就是 MoE「算得少、存得多」的特点，也是后面专家并行要解决的问题。MoE 还引入新的通信：每个 MoE 层前后各一次 AllToAll，把 token 发到专家所在的卡上再收回来。
KV cache（推理时的显存账）
推理分为 prefill（处理整个提示）和逐 token 的生成。prefill 时把每层的 K、V 投影存下来，生成时复用，这就是 KV cache。每条序列的 KV cache 是一个 \([2, S, L, K, H]\) 的数组（2 代表 K 和 V），字节数为：
$$m_\text{KV} = 2\cdot S\cdot L\cdot K\cdot H \cdot \text{bytes}$$Llama 3 8B（\(L=32, K=8, H=128\)）用 BF16 存：每 token \(2\times 32\times 8\times 128\times 2\) 字节 = 128 KiB，一条 8k 序列约 1.07 GB。如果是 MHA（\(K = N = 32\)），会是 4 倍。这就是 GQA/MQA 的动机。推理的完整账见Transformer 推理的数学。
13. MFU 与 HFU
知道了一步需要多少 FLOPs，就能衡量训练「用满」了多少硬件。设一步的时长为 \(t_\text{step}\)，单卡峰值算力为 \(P_\text{peak}\)（FLOP/s），卡数为 \(n\)：
$$\text{MFU} = \frac{\text{模型 FLOPs}/t_\text{step}}{n\cdot P_\text{peak}},\qquad \text{HFU} = \frac{\text{硬件 FLOPs}/t_\text{step}}{n\cdot P_\text{peak}}$$- 硬件 FLOPs（hardware FLOPs）：硬件实际执行的运算，包含重计算。由它得到的利用率叫 HFU（hardware FLOPs utilization）。
- 模型 FLOPs（model FLOPs）：模型前向加反向本身所需的运算，不含重计算。由它得到的叫 MFU（model FLOPs utilization），这个概念来自 PaLM 论文。
为什么更看重 MFU？因为最终关心的是在给定数据上训完要多久。如果一种硬件显存足够大、不需要重算，总 FLOPs 更少却训得更快，它应该被奖励而不是被惩罚。HFU 会奖励「多做无用功」，MFU 不会，所以 MFU 更多反映模型和训练方案本身，便于跨实现比较。
两者的关系：完全重算下，硬件 FLOPs ≈ \(8ND\)，模型 FLOPs = \(6ND\)，于是 \(\text{HFU} \approx \tfrac{4}{3}\,\text{MFU}\)；选择性重算下，比值约为 \(1 + s/(6h)\)。
几个真实数字
- Llama 3 405B 的三种训练配置：8,192 张 H100 上每卡 430 TFLOP/s、MFU 43%；16,384 张卡上 400 TFLOP/s、41%；128k 长上下文阶段 380 TFLOP/s、38%。MFU 的分母是 H100 SXM 的 BF16 稠密峰值 989 TFLOP/s（430/989 ≈ 43.5%）。卡数翻倍后 MFU 下降，是因为全局 batch 保持 16M token 不变，每个数据并行组分到的 batch 变小了。
- DeepSeek-V3：预训练 14.8T token 用了 2.664M H800 GPU 小时（加上长上下文扩展和后训练共 2.788M），激活参数 37B。H800 的计算单元和 H100 SXM 相同（被削减的是 NVLink 带宽），FP8 稠密峰值约 1,979 TFLOP/s。预训练的总可用算力是 \(2.664\text{e}6 \times 3600 \times 1.979\text{e}15 \approx 1.90\text{e}25\) FLOPs，需要的是 \(6 \times 37\text{e}9 \times 14.8\text{e}12 \approx 3.3\text{e}24\)，以 FP8 峰值为分母的利用率约 17%。分母用 FP8 是因为主要的矩阵乘法都用 FP8 做；而注意力等仍用 BF16 的部分既没算进 \(6ND\)，峰值也只有 FP8 的一半，所以硬件的实际繁忙程度比 17% 高。
小心规格表上的数字。NVIDIA 规格表里 Tensor Core 的数字通常是「带结构化稀疏」的，是稠密值的 2 倍。H100 SXM 规格表写 BF16 1,979 TFLOPS 并标了星号，稠密值是 989。算 MFU 时要用稠密峰值，否则 MFU 会被低估一半。
一个好用的换算
把 MFU 换算成每卡吞吐：
$$\text{tokens/s/GPU} = \frac{\text{MFU}\cdot P_\text{peak}}{6N}$$Llama 3 8B 在 H100 上、MFU 40% 时：\(0.4 \times 989\text{e}12 / (6 \times 8.03\text{e}9) \approx 8{,}200\) token/s/GPU。这个数在看训练日志时非常有用：日志报的吞吐乘以 \(6N\) 再除以峰值，就是你的 MFU。（长序列时别忘了加上注意力项，否则 MFU 会被低估。）
14. 交互：训练要多少天
把上面的东西串起来：
$$t_\text{train} = \frac{C_\text{train}}{n \cdot P_\text{peak} \cdot \text{MFU}}$$下面的估算器用 \(6ND\) 加上注意力项（按因果减半），也给出完全重算时的 HFU。峰值都取 BF16 稠密值：H100/H200 989、B200 2,250（HGX B200 的规格写的是 8 卡 36 PFLOPS，这是稀疏值，稠密减半后每卡 2.25 PFLOPS）、A100 312、TPU v5p 459、TPU v5e 197 TFLOP/s。
注意力项按 \(6\cdot T\cdot (a\,d_h)\cdot L\) FLOPs/token 估算（训练、因果减半），其中层数和注意力宽度由参数量按 Llama 3 的形状插值得到。HFU 一栏假设完全重算（多一遍前向）。
默认是 Llama 3 405B 的配置：15.6T token、16,384 张 H100、MFU 41%，只算 \(6ND\) 是约 66 天，加上 8k 序列的注意力项约 69 天。作为对照，Llama 3.1 的模型卡记录 405B 一共用了 30.84M H100 GPU 小时，按 16,384 张卡均摊是约 1,880 小时、78 天。反过来算平均利用率：\(3.8\text{e}25 / (30.84\text{e}6 \times 3600 \times 989\text{e}12) \approx 35\%\)，比上一节稳态运行时的 38–43% 低。差距来自稳态之外的时间：GPU 小时是整个训练的总账，包括 MFU 本来就更低的长上下文阶段，也包括故障带来的中断和重启。Llama 3 预训练期间，54 天里发生了 466 次作业中断，有效训练时间约占 90%。
15. 小结：一张速查表
| 量 | 公式 | 备注 |
|---|---|---|
| 参数量（Llama 式） | \(L(2h\,ad_h + 2h\,h_{kv} + 3hF + 2h) + 2Vh\) | 粗略 \(\approx 12Lh^2\)（GPT 式） |
| 静态显存 | \(16N\)（或 \(20N\)）字节 | BF16 混合 + FP32 主权重 + Adam |
| 激活（GPT 结构） | \(L\,sbh(34 + 5as/h)\) 字节 | 含 dropout 掩码、朴素注意力 |
| 激活（Llama 式） | \(L\cdot 2sb(6h + 2h_{kv} + 4F)\) 字节 | FlashAttention、无 dropout |
| 完全重算后 | \(L\cdot 2sbh\) + 一层工作集 | FLOPs \(6ND\to 8ND\) |
| 全局 batch | \(gbs = mbs \times grad\_acc\) | 数据并行后再乘 \(dp\) |
| 矩阵乘法 | \(2mnk\)；训练 \(6mnk\) | 前向 : 反向 = 1 : 2 |
| 训练算力 | \(6ND\) + \(12BT^2NH\cdot L\) | 注意力在 \(T > 8D\) 后占主导 |
| MoE | \(6N_\text{active}D\) | 显存按总参数 |
| KV cache | \(2SLKH\) × 字节 | 每条序列 |
| MFU | \(\frac{6N\cdot \text{tokens/s}}{n\,P_\text{peak}}\) | 用稠密峰值 |
回到开头的问题：Llama 3 8B 的 16 字节/参数就已经是 128 GB，激活值还会再加几十 GB——单卡装不下，这是必须多卡训练的直接原因。但切模型之前还有一个问题：给定的一次运算，瓶颈到底在算力、显存带宽还是网络上？Roofline 回答的就是这个。
参考与延伸阅读
这一篇的显存部分主要整理自 Ultra-Scale Playbook 的前两章，FLOPs 部分主要整理自 How to Scale Your Model 的第 4 章。想看真实的显存 profiling 曲线和更多实验，推荐直接读 Playbook 的「First Steps: Training on One GPU」；第 4 章末尾的习题也很值得做一遍。
- Nouamane Tazi 等，The Ultra-Scale Playbook: Training LLMs on GPU Clusters：High-Level Overview；First Steps: Training on One GPU（Memory usage in transformers、Profiling the memory usage、Memory for weights/grads/optimizer states、Memory for activations、Activation recomputation、Gradient accumulation）；附录 A2 Typical Scales in LLM Training。
- Jacob Austin 等，How to Scale Your Model, Part 4: All the Transformer Math You Need to Know：Counting Dots、Forward and reverse FLOPs、Global FLOPs and Params Calculation、Sparsity and Mixture-of-Experts、Gradient checkpointing、Key-Value caching 及习题。
- Vijay Korthikanti 等，Reducing Activation Recomputation in Large Transformer Models（2022）：激活显存公式、选择性重算、附录 A 的 FLOPs 公式。
- Llama Team, Meta，The Llama 3 Herd of Models（2024）：表 3（模型配置）、表 4（并行配置与 MFU）、预训练 batch 调度、\(3.8\times10^{25}\) FLOPs、训练中断统计。
- Llama 3.1 Model Card：各尺寸的 H100 GPU 小时数；Hugging Face 上的 Llama 3 config：模型维度与词表。
- DeepSeek-AI，DeepSeek-V3 Technical Report（2024）：batch size 调度、RMSNorm 与 MLA 上投影的重算、表 1 的 GPU 小时。
- Tianqi Chen 等，Training Deep Nets with Sublinear Memory Cost（2016）。
- NVIDIA H100 与 HGX B200 规格页。