薛小嵩blog

MFU估算

 
  • 2025年 09月28日
  • 0 阅读

1. 浮点计算基础概念

1.1 FLOP（Floating Point Operation）

FLOP 表示一次基本的浮点运算，例如加法 $x + y$ 或乘法 $x \times y$。

1.2 FLOPs

FLOPs 表示“总共做了多少浮点运算”，即计算工作量的度量。

例如，训练 GPT-3 的总计算量约为：

\[3.14 \times 10^{23} \text{ FLOPs}\]

1.3 FLOP/s 或 FLOPS

FLOP/s（FLOPS） 表示“每秒能执行多少次浮点运算”，是硬件的速度指标。

例如：

  • NVIDIA A100：约 312 TFLOP/s
  • NVIDIA H100：约 1979 TFLOP/s（带稀疏性） ≈ 989 TFLOP/s（实际）

2. 线性模型的 FLOPs 计算

2.1 示例：线性层矩阵乘法

以一个简单线性层为例：

y = x @ w

设：

  • $B$ = batch size（样本数）
  • $D$ = 输入维度
  • $K$ = 输出维度

矩阵乘法 $(B \times D) \times (D \times K)$ 包含：

  • 1 次乘法：$x[i][j] \times w[j][k]$
  • 1 次加法

因此，实际 FLOPs 为：

\[\text{actual_num_flops} = 2 \times B \times D \times K\]

2.2 其他算子的 FLOPs 量级

矩阵乘法是所有神经网络中 最主要的 FLOPs 来源。

其他算子的 FLOPs 数量级：

算子类型 FLOPs 复杂度 说明
ReLU、GELU、SiLU $O(B \times D)$ 每个元素只做 1~2 次运算
LayerNorm、Softmax、Dropout $O(B \times D)$ 或 $O(B \times L)$ 序列级操作，开销较小

结论： 在 FLOPs 统计中，除矩阵乘法外的算子几乎可忽略。

3. FLOPs 与Transformer

3.1 Transformer 模型 FLOPs 分解

记：

  • $B$：batch size
  • $L$：sequence length（序列长度）
  • $H$：hidden size（隐藏维度）

各模块 FLOPs 近似如下：

模块 计算形式 近似 FLOPs 含义
Attention QKV projection $[B, L, H] \times [H, 3H]$ $2BLH(3H)$ 从 hidden 生成 Q、K、V
Attention 乘积 $(QK^\top)$ $[B, L, H] \times [B, H, L]$ $2BL^2H$ 序列内自注意力
Attention 输出 $(AV)$ $[B, L, L] \times [B, L, H]$ $2BL^2H$ 加权聚合结果
FFN 前后两层线性 $[B, L, H] \times [H, 4H] + [B, L, 4H] \times [4H, H]$ $2BLH(8H)$ MLP 两个线性层

综合可得：

\[\text{FLOPs per token} \approx 6H^2\]

（包含 attention + MLP 的主要矩阵乘法）

3.2 总体 FLOPs 估算公式

整个训练的总 FLOPs 近似为：

\[\text{Total FLOPs} \approx 6 \times N_{\text{params}} \times N_{\text{tokens}}\]

其中：

  • $N_{\text{params}}$：模型参数量（≈ $O(H^2)$）
  • $N_{\text{tokens}}$：训练总 token 数
  • 系数 6：由每个 Transformer Block 的 Attention + MLP 的加权 FLOPs 得出。

计算展开：

输入 X: [B, L, H]
 ├─ QKV: 3 × (2 × B × L × H × H)
 ├─ Attention(QKᵀ + AV): 4 × (2 × B × L² × H)
 └─ MLP: 8 × (2 × B × L × H × H)
 ≈ O(6 × B × L × H²)

4. 理论算力与实际速度

4.1 代码计时与 MFU

实际测量矩阵乘法性能：

actual_time = time_matmul(x, w)
actual_flop_per_sec = actual_num_flops / actual_time

与硬件规格表比较：

promised_flop_per_sec = get_promised_flop_per_sec(device, x.dtype)
mfu = actual_flop_per_sec / promised_flop_per_sec

得到 MFU (Model FLOPs Utilization)：

\[\text{MFU} = \frac{\text{实际算力}}{\text{理论峰值算力}}\]

一般来说：

  • MFU ≥ 0.5（即 50%）已属优秀。
  • 余下约 50% 的损耗主要来自 kernel 启动、通信延迟、内存带宽限制 等非计算部分。

5. 前向与反向 FLOPs 分析

5.1 模型与张量形状

以两层线性模型为例：

x --w1--> h1 --w2--> h2 -> loss

维度设定：

  • $x$: [B, D]
  • $w_1$: [D, D]
  • $w_2$: [D, K]
  • $h_1 = x @ w_1$: [B, D]
  • $h_2 = h_1 @ w_2$: [B, K]
  • $loss = mean(h_2^2)$

5.2 前向 FLOPs

矩阵乘法 $(m \times n) @ (n \times p)$ 的 FLOPs 约为 $2mnp$。

  • 第一层：$2BD^2$
  • 第二层：$2BDK$

合计前向 FLOPs：

\[\text{Forward FLOPs} = 2B(D^2 + DK)\]

5.3 反向 FLOPs（链式法则）

针对 $w_2$

  1. 计算 $w_2$ 梯度： $(D \times B) @ (B \times K)$ → $2BDK$
  2. 计算 $h_1$ 梯度： $(B \times K) @ (K \times D)$ → $2BDK$

合计： 反向 FLOPs（与 $w_2$ 相关）= $4BDK$

针对 $w_1$

  1. 计算 $w_1$ 梯度： $(D \times B) @ (B \times D)$ → $2BD^2$
  2. 若需 $x.grad$： $(B \times D) @ (D \times D)$ → $2BD^2$

合计： 反向 FLOPs（与 $w_1$ 相关）= $4BD^2$

若不需要 $x.grad$，可省略最后一项。

5.4 加总与“2/4/6 规则”

综合上述结果（包含 $x.grad$）：

阶段 FLOPs 说明
前向 $2 \times$ 数据点 × 参数数 单次推理
反向 $4 \times$ 数据点 × 参数数 梯度计算与传播
总计 $6 \times$ 数据点 × 参数数 一次训练迭代

最终得到著名的一阶估算公式：

\[\text{训练总 FLOPs} \approx 6 \times N_{\text{params}} \times N_{\text{tokens}}\]

该经验公式来源于线性层推导，在 Transformer、LLaMA、GPT 等架构中仍然适用，因为主导 FLOPs 的部分始终是矩阵乘法。

参考资料： The FLOPs calculus of language model training — Dzmitry Bahdanau, Medium

搜索