参数化估算训练阶段的 FLOPs、访存量与算术强度

你好。上一课区分了 MoE 的总参数量、单 token 激活参数量和实际 FLOPs：Top-16 / 384 只让 token 使用 $1/24$ 的专家参数，并不意味着完整训练成本就是总参数量的 $1/24$ 。

这一课把这个区分变成可用于性能建模的账本。目标不是背下某个“每 token 几 FLOPs”的常数，而是能够对任意候选模型形状、microbatch 和路由分布，分别估算：

前向、反向、重算与优化器更新的 FLOPs；
HBM 读写的下界与实现修正项；
MoE token 分发的网络字节数；
算术强度（Arithmetic Intensity, AI），从而判断低带宽 ASIC 更可能卡在计算还是数据搬运。

后续讨论 Roofline、关键算子需求和 5T 模型的容量规划，都会直接使用这套符号。

先定义性能账本的变量

为避免把“序列长度”“层数”“token 数”混为一谈，以下记号在本课程后续沿用：

符号	含义
$B$	每个数据并行 rank 的 microbatch 中序列条数
$S$	每条序列的 token 长度
$M = B S$	一个 rank 在一次前向中处理的 token 数
$L_{layer}$	Transformer 层数
$D$	模型隐藏维度
$F$	FFN 中间维度
$n_{h}$	注意力 query 头数
$H$	每个头的维度，满足 $D = n_{h} H$
$E$	每个 MoE 层的专家数，目标为 $384$
$k$	每个 token 选择的专家数，目标为 $16$
$b$	一个元素在 HBM 中占用的字节数；BF16 时通常为 $2$
$a_{i}$	一个 MoE 层中被路由到专家 $i$ 的 token-expert assignment 数
$p_{i}$	专家 $i$ 实际送入 GEMM 的行数，可能因 padding 而大于有效 assignment 数

对一个没有丢弃 assignment 的 MoE 层：

i = 1 \sum E a_{i} = M k

这很重要：一个 token 选中 $k = 16$ 个专家，会形成 16 个 token-expert assignment。因此，MoE 专家 GEMM 的“有效 token 行数”不是 $M$ ，而是 $M k$ 。

先花一点时间看一份公式参考。它把矩阵乘、Transformer、MoE 和算术强度放在同一套记号下；不必试图记住所有表格，但应理解“训练 FLOPs 约为前向 GEMM FLOPs 的三倍”以及 MoE 的 batch 被专家数稀释这两个结论。

All the Transformer Math You Need to Know | How To Scale Your Model

阅读 Google DeepMind 等作者的《All the Transformer Math You Need to Know》。它提供了本课所需的矩阵乘训练 FLOPs、注意力与 FFN 分解、MoE 稀疏性以及专家 GEMM 算术强度公式。

先读 “Forward and reverse FLOPs”：从前反向计数理解一个大 GEMM 的训练计算为何约为前向的三倍。接着读 “Global FLOPs and Params Calculation” 下的 “MLPs” 与 “Attention” 两个小节。重点核对门控 MLP 的三次投影，以及 Q、K、V、O 投影和两次注意力矩阵乘各自如何随 B、T、D、F 增长。然后在 “Miscellaneous Math” 的 “Sparsity and Mixture-of-Experts” 小节，阅读 MoE 通信说明，注意一个 MoE 层有两次 token 重排通信。最后到 “A Few Problems to Work” 的 Question 8，读专家算术强度推导。把其中的结论理解为：专家数越多、每 token 选中的专家比例越低，就越需要更大的 token batch 来复用专家权重。

从一个 GEMM 开始：FLOPs 和 HBM 字节不是同一件事

几乎所有 Transformer 主计算都可以拆成：

Y = X W

其中：

X \in R^{M \times D}, W \in R^{D \times F}, Y \in R^{M \times F}

这里 $M = B S$ 是当前 rank 的 token 数。每个输出元素需要约 $D$ 次乘法与 $D$ 次累加，所以前向计算量为：

F_{fwd} \approx 2 M D F

反向传播有两项主要矩阵乘：

\nabla_{X} = \nabla_{Y} W^{T}

\nabla_{W} = X^{T} \nabla_{Y}

它们各自约为：

2 M D F

因此，一个训练中的线性层可按下式估算：

F_{train, linear} \approx 2 M D F + 2 M D F + 2 M D F = 6 M D F

这就是常见的“训练约为前向三倍计算量”的来源。注意，这个比例对大矩阵乘很有用；它不是说训练阶段每一种操作、每一种数据移动都恰好是前向的三倍。

对于前向 HBM 流量，最乐观的下界是：

V_{fwd, linear} ≳ b (M D + D F + M F)

三项分别对应：

读取输入 $X$ ；
读取权重 $W$ ；
写出输出 $Y$ 。

这里的 $≳$ 很关键。它表示下界，而不是 profiler 必然看到的真实字节数。真实系统可能因为缓存容量不足、分块重读、未融合中间结果、转置布局或动态 shape 而移动更多数据。

于是，该 GEMM 的前向算术强度为：

I_{linear} = \frac{F _{fwd, linear}}{V _{fwd, linear}} \approx \frac{2 M D F}{b ( M D + D F + M F )}

单位是 FLOPs/byte。

两个极限情况

当 $M$ 很小，权重读取 $D F$ 主导 HBM 流量：

I_{linear} \approx \frac{2 M D F}{b D F} = \frac{2 M}{b}

对于 BF16 权重， $b = 2$ ，因此：

I_{linear} \approx M

也就是说，小 token batch 时，GEMM 的算术强度大致只由 token 数决定。权重几乎没有被复用，低带宽 ASIC 很容易被 HBM 限制。

当 $M$ 足够大，激活读写主导：

I_{linear} \approx \frac{2 D F}{b ( D + F )}

继续增大 batch 仍可提高吞吐，但不会无限提高算术强度，因为此时需要持续搬运输入和输出激活。

用 Roofline 把 FLOPs 与字节变成时间下界

仅有 FLOPs 不能说明一个算子会运行多快。设备有至少两条基本上限：

峰值矩阵计算能力 $P_{peak}$ ，单位 FLOPs/s；
可达 HBM 带宽 $B W_{HBM}$ ，单位 bytes/s。

对某算子，其理论时间下界为：

t_{op} \geq max (\frac{F _{op}}{P _{peak}}, \frac{V _{op}}{B W _{HBM}})

定义硬件的临界算术强度：

I^{⋆} = \frac{P _{peak}}{B W _{HBM}}

则有一个简洁的判断规则：

I_{op} < I^{⋆} \Rightarrow HBM 带宽受限

I_{op} > I^{⋆} \Rightarrow 计算受限的可能性较高

这里说“可能性较高”，不是“必然计算受限”。低效 tile、碎片化专家 GEMM、kernel launch 开销、通信等待和负载倾斜，都可能使实际性能远低于 Roofline 上限。

对于你的目标 ASIC，“大算力、低带宽”意味着：

I^{⋆}

可能较高。因此，仅仅看到某算子有大量 FLOPs 并不够；还必须问：这些 FLOPs 是否能以足够高的权重复用率完成？

一个 decoder 层的前向 FLOPs

先用标准多头注意力和门控 FFN 建立基础公式，再把稠密 FFN 换成 MoE。

注意力投影

标准注意力的 Q、K、V、O 投影可以近似视为四个：

(M \times D) (D \times D)

矩阵乘。因此前向 FLOPs 为：

F_{attn, proj, fwd} \approx 4 \cdot 2 M D^{2} = 8 M D^{2}

训练 FLOPs 约为：

F_{attn, proj, train} \approx 24 M D^{2}

如果采用 GQA，K、V 的投影维度小于 $D$ ，应将这四个 $D^{2}$ 项替换为实际 Q、K、V、O 投影矩阵的参数量；不要无条件套用标准 MHA 的 $8 M D^{2}$ 。

注意力的长度平方项

注意力还有两个不由参数量直接决定的矩阵乘：

Q K^{T}

以及：

P V

对全长序列，前向近似为：

F_{attn, quad, fwd} \approx 4 B S^{2} D

训练近似为：

F_{attn, quad, train} \approx 12 B S^{2} D

对 causal attention，理论上三角区域只包含约一半的位置，因此“有用 FLOPs”大约可减半。但实际执行是否达到这一节省，取决于 attention kernel 是否真正利用因果结构，而不是构造并计算完整方阵。

稠密门控 FFN

SwiGLU 一类门控 FFN 典型包含三次投影：

X W_{up}, X W_{gate}, Z W_{down}

每个主要投影的前向成本约为 $2 M D F$ ，所以：

F_{FFN, fwd} \approx 6 M D F

训练阶段：

F_{FFN, train} \approx 18 M D F

若 $F \approx 4 D$ ，则单层门控 FFN 前向约为：

F_{FFN, fwd} \approx 24 M D^{2}

这解释了为什么在普通上下文长度下，FFN 往往主导 Transformer 的参数量和主矩阵计算。

将稠密 FFN 替换为 Top- $k$ MoE

在 MoE 层，专家 $i$ 接到 $a_{i}$ 个 assignment。一个 gated expert FFN 的前向成本近似为：

6 a_{i} D F

对全部专家求和：

F_{expert, fwd, useful} \approx 6 D F i = 1 \sum E a_{i}

若没有 token 丢弃：

i = 1 \sum E a_{i} = M k

因此：

F_{expert, fwd, useful} \approx 6 M k D F

训练成本近似为：

F_{expert, train, useful} \approx 18 M k D F

这反映了一个容易遗漏的事实：

对于单个 token，Top-16 意味着它要经历 16 个专家 FFN 的计算，而不是一个专家 FFN 的计算。

MoE 的稀疏性节省来自“不计算其余 $E - k$ 个专家”，不是说 Top- $k$ 的计算等价于单个稠密 FFN。

路由器本身通常包含一个：

(M \times D) (D \times E)

投影，前向成本为：

F_{router, fwd} \approx 2 M D E

训练成本约为：

F_{router, train} \approx 6 M D E

Top- $k$ 选择、排序、prefix sum 和 token 重排的 FLOPs 通常不是主导项，但它们可能具有很差的访存模式、动态 shape 和同步开销。因此，不能因为 router FLOPs 小，就把路由路径视为“免费”。

padding 会使“实际执行 FLOPs”超过有用 FLOPs

上一课已区分算法层面的有效计算与硬件实际执行。现在可将差异写成公式。

设每个专家的容量为：

C = ⌈ c \frac{M k}{E} ⌉

其中 $c$ 是容量因子。若每个专家都被固定形状地执行为 $C$ 行，即使其真实收到的 assignment 更少，也会产生：

i = 1 \sum E p_{i} = E C

个实际 GEMM 行。

此时专家前向执行 FLOPs 为：

F_{expert, fwd, executed} \approx 6 D F i = 1 \sum E p_{i}

定义 padding 放大因子：

γ_{pad} = \frac{\sum _{i} p _{i}}{\sum _{i} a _{i}}

则：

F_{expert, fwd, executed} = γ_{pad} F_{expert, fwd, useful}

当使用固定容量，且没有严重溢出丢弃时：

γ_{pad} \approx c

但这只是理想化近似。真实 $γ_{pad}$ 还取决于负载倾斜、容量截断策略、空专家是否启动 GEMM，以及 Grouped GEMM 是否能够按实际 token 数动态打包。

该图展示了同一 GPU 上的 token 先按所选专家聚集，再由一次融合的 GroupGEMM 调用执行多个专家 FFN；其目标是让专家权重在共享存储中被多个 token 复用，而不是为每个 token 单独启动小矩阵乘。

对性能建模，应同时记录：

F_{useful}

和：

F_{executed}

前者用于描述算法工作量；后者更接近芯片实际做了多少矩阵计算。二者之差并不表示模型“学到了更多”，而是 padding、静态容量和布局代价。

MoE 专家 GEMM 为什么在低带宽 ASIC 上困难

对某个活跃专家 $i$ ，以其 token 数 $p_{i}$ 近似其 GEMM 行数。忽略中间激活的额外读写，一个门控 expert FFN 的粗略前向模型为：

F_{i, fwd} \approx 6 p_{i} D F

理想 HBM 字节量则至少包括：

读取三组 FFN 权重，约 $3 b_{w} D F$ ；
读取和写出激活，量级与 $p_{i} D$ 、 $p_{i} F$ 成正比。

因此，专家 $i$ 的算术强度可粗略写为：

I_{i, expert} \approx \frac{6 p _{i} D F}{3 b _{w} D F + b _{a} ( α p _{i} D + β p _{i} F )}

其中 $α$ 、 $β$ 是实现相关常数：它们会随融合程度、是否物化中间激活、数据格式和缓存策略变化。

若 $p_{i}$ 很小，权重读取主导分母，则：

I_{i, expert} \approx \frac{2 p _{i}}{b _{w}}

对 BF16 权重：

b_{w} = 2

于是：

I_{i, expert} \approx p_{i}

这与普通线性层的“小 batch”结论相同：一个专家处理的 token 数越少，权重复用越差。

在完全均衡的理想路由下：

p ˉ = \frac{M k}{E}

对于目标配置：

p ˉ = \frac{16 M}{384} = \frac{M}{24}

也就是说，即使整个 rank 处理 $M$ 个 token，平均每个专家只得到 $M /24$ 个 assignment。若 $M = 2, 400$ ，平均每个专家也只有：

p ˉ = 100

行。对“高算力、低带宽”的 ASIC，这可能远不足以让专家权重获得充分复用。

因此，MoE 调优不能只问“全局 batch 有多大”，还要问：

\frac{M k}{E}

以及实际路由下的：

i min p_{i}, median (p_{i}), i max p_{i}

这些值共同决定 Grouped GEMM 的形状、尾部时延和有效算术强度。

128k 长上下文：注意力 FLOPs 与 HBM 流量

对于标准注意力，前向的二次项是：

F_{attn, quad, fwd} \approx 4 B S^{2} D

它与序列长度 $S$ 的平方成正比。另一方面，投影和 FFN 主计算通常仅与 token 总数 $M = B S$ 成正比。

以普通稠密门控 FFN 且 $F \approx 4 D$ 为例，单层前向中：

F_{linear, fwd} \approx 8 M D^{2} + 24 M D^{2} = 32 M D^{2}

注意力二次项相对线性项的比例为：

\frac{F _{attn, quad, fwd}}{F _{linear, fwd}} \approx \frac{4 B S ^{2} D}{32 B S D ^{2}} = \frac{S}{8 D}

因此，当：

S ≳ 8 D

注意力二次计算会开始达到或超过投影与 FFN 的线性计算量。

对 MoE，这个交叉点不能直接照搬，因为专家 FFN 的计算应使用：

6 M k D F

而不是稠密 FFN 的：

6 M D F

此外，专家的 $F$ 是否缩小、哪些层是 MoE、注意力是否采用 GQA、是否采用上下文并行，都会改变结果。正确做法是把每一项代入实际模型配置，而不是从“总参数 5T”直接推导注意力占比。

FlashAttention 改变的是字节量，不是二次 FLOPs

朴素注意力若将 score 矩阵和 softmax 概率写入 HBM，会引入量级为：

O (B n_{h} S^{2})

的额外 HBM 流量。FlashAttention 类实现通过分块，在片上保存中间 softmax 状态，尽量避免物化完整注意力矩阵。

理想化的 FlashAttention 前向 HBM 下界接近：

V_{attn, flash} ≳ 4 b B S D

而不是朴素实现中的：

O (b B n_{h} S^{2})

因此，FlashAttention 不会消除：

O (B S^{2} D)

的计算，但能显著改变注意力的算术强度和 HBM 瓶颈位置。这对 $128 k$ 训练是算子库的硬要求，而非可有可无的优化。

前向、反向、重算、优化器：一整个训练 step 的公式

设第 $ℓ$ 层的前向主计算为：

F_{ℓ, fwd} = F_{ℓ, attn, proj} + F_{ℓ, attn, quad} + F_{ℓ, router} + F_{ℓ, expert} + F_{ℓ, other}

其中可用一阶近似：

F_{ℓ, fwd} \approx 8 M D^{2} + 4 B S^{2} D + 2 M D E + 6 D F i = 1 \sum E p_{ℓ i}

若是非 MoE 的稠密 FFN 层，则将专家项替换为：

6 M D F

对于由大矩阵乘主导的层，反向主计算近似为：

F_{ℓ, bwd} \approx 2 F_{ℓ, fwd, matmul}

因此，训练中的前向加反向可写为：

F_{fwd + bwd} \approx 3 ℓ = 1 \sum L_{layer} F_{ℓ, fwd, matmul}

若有 activation checkpointing 或 rematerialization，还要额外加上被重算部分的前向 FLOPs：

F_{step} \approx F_{fwd + bwd} + F_{remat} + F_{optimizer}

其中：

F_{remat} = r F_{recomputed, fwd}

$r$ 不是固定常数：它取决于保存哪些激活、保存 attention 的哪些状态、是否重算 MoE 路由、以及编译器如何融合图。

优化器 FLOPs 很小，但访存不可忽略

以 AdamW 为例，每次更新会触及所有可训练参数，而不仅是本 step 被选中的专家：

N_{update} = N_{total}

这意味着：尽管 Top-16 降低了专家前反向的激活计算，优化器更新仍要遍历全部 $5 T$ 参数及其状态。

令：

$b_{g}$ ：梯度字节数；
$b_{p}$ ：FP32 master weight 的字节数；
$b_{m}$ 、 $b_{v}$ ：一阶和二阶矩状态的字节数；
$b_{w}$ ：更新后训练权重副本的字节数。

一次未充分融合的 AdamW 更新的 HBM 流量下界可参数化为：

V_{opt} ≳ N_{update} (b_{g} + 2 b_{p} + 2 b_{m} + 2 b_{v} + b_{w})

“ $2$ ”表示读一次、写一次。

例如，一个常见配置是：

BF16 梯度： $b_{g} = 2$ ；
FP32 master weight： $b_{p} = 4$ ；
FP32 一阶矩： $b_{m} = 4$ ；
FP32 二阶矩： $b_{v} = 4$ ；
BF16 更新后模型权重： $b_{w} = 2$ 。

则每个参数的最低访问量级约为：

2 + 2 \cdot 4 + 2 \cdot 4 + 2 \cdot 4 + 2 = 28 bytes

实际融合优化器、参数分片和缓存会改变这个数，但不会改变其本质：

优化器阶段通常 FLOPs 占比低，却可能是一次大规模、低算术强度的全参数流式访问。

这也是 5T MoE 系统中“前反向很稀疏，但 optimizer 仍很重”的具体含义。

HBM 字节账本：不要把“训练 FLOPs 三倍”误用于访存

对一个完整训练 step，建议采用下面的结构，而不是试图用单一倍数估算总字节量：

V_{HBM, step} = V_{fwd} + V_{saved} + V_{bwd} + V_{remat} + V_{opt}

各项含义如下：

阶段	主要 HBM 流量	是否可简单按前向倍数估算
前向	权重读取、输入读取、输出与中间激活写入	只能作为下界
保存激活	为反向保留的 residual、QKV、FFN 中间量、路由信息等	强依赖 checkpoint 策略
反向	读取保存激活与权重，写梯度和中间梯度	不应机械地写为前向两倍
重算	重新读取权重和输入激活，计算被丢弃的前向节点	取决于重算边界
优化器	全参数梯度、主权重、动量、方差和训练权重副本	与总参数相关，不与激活专家数成正比

对算子库团队而言，需要的不是笼统的“HBM 利用率”，而是能回答以下问题的 profile：

attention 是否写出了 $S^{2}$ 大小的中间矩阵；
SwiGLU、RMSNorm、残差与量化是否被融合，避免中间 HBM 往返；
Grouped GEMM 是否因碎片化专家 batch 而重复读取权重；
backward 是否重算了本来可以保存的高代价节点，或保存了本不值得保存的大激活；
优化器是否以大而连续的 shard 更新，还是被大量小 kernel 和额外格式转换切碎。

MoE 的网络流量要与 HBM 流量分开建模

MoE 层通常至少包含两次 All-to-All：

将 token hidden states 分发到专家所在 rank；
将专家输出返回 token 原始归属 rank 并做加权合并。

若 EP group 内总共有 $M_{EP}$ 个原始 token，且每个 token 选择 $k$ 个专家，则一次分发的有效 payload 量级为：

V_{dispatch} \approx M_{EP} k D b_{comm}

一次返回的量级近似相同，因此每层双向 payload 为：

V_{MoE, A2A} \approx 2 M_{EP} k D b_{comm}

若定长 whole-group All-to-All 要求每个目的 rank 发送固定槽位，并按容量 $C$ 填充，则实际通信可能更接近：

V_{MoE, A2A, padded} \approx 2 E C D b_{comm}

这部分是网络字节，不要与 HBM 读写相加后直接称为“访存量”。二者经过不同硬件通道，也具有不同带宽、拓扑和重叠机会。

对于后续的性能模型，至少应分别保留：

V_{HBM}

V_{network}

F_{executed}

只有这样，才能解释一种常见现象：虽然 expert GEMM 的算术强度足够高，但 step 仍然被 All-to-All、token 重排或容量填充拖慢。

应放入 5T MoE 性能表格的最小公式集

对每个候选模型配置、每种 microbatch 和每个并行布局，建议至少填入以下字段。

项目	参数化估算
每 rank token 数	$M = B S$
每 MoE 层有效 assignment 数	$\sum_{i} a_{i} = M k$
平均每专家 token 数	$p ˉ = M k / E$
Attention projection 前向 FLOPs	$8 M D^{2}$ ，GQA 时按实际投影矩阵替换
Attention 二次前向 FLOPs	$4 B S^{2} D$ ，causal kernel 可按有效三角区域调整
Router 前向 FLOPs	$2 M D E$
MoE FFN 有效前向 FLOPs	$6 D F \sum_{i} a_{i}$
MoE FFN 实际前向 FLOPs	$6 D F \sum_{i} p_{i}$
大 GEMM 训练 FLOPs	约为相应前向 matmul FLOPs 的 $3$ 倍
重算 FLOPs	$r F_{recomputed, fwd}$
优化器状态访问	$V_{opt} ≳ N_{update} (b_{g} + 2 b_{p} + 2 b_{m} + 2 b_{v} + b_{w})$
MoE All-to-All payload	$2 M_{EP} k D b_{comm}$
算术强度	$I = F / V_{HBM}$
Roofline 临界点	$I^{⋆} = P_{peak} / B W_{HBM}$

这里最重要的纪律是：每一项都要附带它的范围和假设。例如：

是 BF16 还是 FP8 权重；
$M$ 是单 rank、单 EP group 还是全局 batch；
$p_{i}$ 是有效 token 数还是 padding 后的 GEMM 行数；
attention 是朴素实现还是 FlashAttention；
网络字节是有效 payload 还是定长通信的真实 buffer 大小；
优化器状态是否全 FP32，是否采用分片。

没有这些假设的 FLOPs 数字很容易精确到小数点，却没有工程决策价值。

小结

本课建立了用于训练性能建模的三本账：

计算账本

F_{train, linear} \approx 6 M D F

大矩阵乘的训练 FLOPs 约为前向的三倍；完整 Transformer 还必须显式加上注意力二次项、router、MoE expert、重算和优化器更新。

HBM 账本

V_{HBM, step} = V_{fwd} + V_{saved} + V_{bwd} + V_{remat} + V_{opt}

FLOPs 与字节不存在统一的“三倍规则”。尤其对 5T 模型，优化器会流式触及所有参数与状态，即使前反向只激活 Top-16 专家。

MoE 稀疏执行账本

i \sum a_{i} = M k

p ˉ = \frac{M k}{E}

F_{expert, fwd, executed} \approx 6 D F i \sum p_{i}

专家数越多、每专家 token 越少，权重复用越差；padding、碎片化与专家倾斜会使实际执行 FLOPs 和数据移动偏离“理想激活参数量”推导。

下一课将把这些单设备阶段模型接到系统指标上：吞吐、时延、扩展效率、MFU、HFU 和通信占比，并讨论为什么一个低 MFU 的训练 step 不一定是“矩阵算力不够用”。

Recommended resources