第 07 章　参数、FLOPs、显存与吞吐怎样估算

上一章已经把 GPT 搭了起来。接下来很自然会问：这个模型要算多久，需要多少内存？只知道参数数量还不够。“模型多大”可能指参数数量、权重文件字节数、一次计算的浮点运算量，也可能指运行时峰值内存。

这些量会相互影响，但回答的是不同问题。比如模型参数完全不变，只把输入序列加长一倍，显式保存的注意力分数矩阵就可能扩大到四倍。本章把这些成本分开计算，再看预算能帮助我们做哪些判断。

上游性能 notebook 使用 thop 统计运算量，搜索设备能容纳的批大小，并估算 GPU 的模型 FLOPs 利用率。我们会对照它的做法解释计数和测量的边界；本章没有运行其中的大型模型、内存不足（OOM）搜索或 GPU 性能测试。固定版本性能分析

7.1 从一个矩阵乘法开始记账

先看最常见的一步：矩阵 $A \in R^{m \times k}$ 乘 $B \in R^{k \times n}$ 。结果有 $m n$ 个元素；算出每个元素，需要 $k$ 次乘法和约 $k$ 次加法，所以总量常近似写成 $2 m k n$ FLOPs。

有些工具报告的是 MACs，也就是乘加对数。按这里一次乘法、一次加法各算一次运算的约定，MACs 要乘 2 才得到 FLOPs。严格说，加法总数是 $m n (k - 1)$ ；做大矩阵预算时，通常忽略这点差异。

这里还要分清两个很像的写法：FLOPs 表示运算总数，FLOP/s 表示每秒运算率。前者告诉我们一次批量前向有多少工作，后者告诉我们硬件每秒能做多少工作。后面估算吞吐时，需要把两者放在各自的位置上。

7.2 一个普通 Transformer block 的主要运算

有了矩阵乘法的计数方法，就可以逐项拆开一个 Transformer block。沿用前文的符号： $B$ 是批量， $T$ 是序列长度， $D$ 是模型宽度， $H$ 是头数，每头宽度 $d = D / H$ 。这里先只算主要矩阵乘法，暂时忽略偏置、归一化、激活和 softmax 等项：

部分	运算量近似
Q、K、V 投影	$6 B T D^{2}$
注意力输出投影	$2 B T D^{2}$
两个 FFN 投影，扩展到 $4 D$	$16 B T D^{2}$
$Q K^{⊤}$	$2 B H T^{2} d = 2 B T^{2} D$
$A V$	$2 B T^{2} D$

把表中各项加起来，每层约为 $24 B T D^{2} + 4 B T^{2} D$ 。模型有 $L$ 层，这部分就再乘 $L$ 。若输出头还要为每个位置计算全词表 logits，则另加约 $2 B T D V$ ，其中 V 是词表大小。embedding 查表主要涉及读取内存，不能直接套用密集矩阵乘法的计数方式。

因果掩码虽然只允许访问一半左右的位置，朴素的稠密（dense）矩阵乘法仍可能先算出整个方阵，再把不允许的位置遮住。因此，mask 呈三角形，并不意味着计算量自动减半。只有具体内核确实跳过了相应的块，才能按三角结构折减实际运算量。

对本书 $B = 2, T = 6, D = 16, L = 2, V = 16$ ，主矩阵前向预算为

2 (24 \times 2 \times 6 \times 16^{2} + 4 \times 2 \times 6^{2} \times 16) + 2 \times 2 \times 6 \times 16 \times 16 = 162816.

这个数来自上面的公式推导。它没有计入 softmax、LayerNorm、激活和内存搬运，也没有经过 profiler 测量，因此不能直接换成某台 CPU 上的运行时间。

7.3 为何反向常粗估为前向的两倍

前面算的是前向。训练还要反向传播，这部分为什么常按前向的两倍估算？对 $Y = X W$ ，反向需要计算 $d X = d Y W^{⊤}$ 与 $d W = X^{⊤} d Y$ ，也就是再做两次与前向同量级的矩阵乘法。当这些乘法占主要成本时，反向可粗估为前向的两倍，前后向合计约三倍。

使用这个三倍系数时，要看实际执行了什么。冻结参数后，部分权重梯度可能不必计算；为了节省内存而重算激活，又会增加前向工作。优化器更新、通信和数据加载也都没有包含在内。

常见的“每 token 约 $6 N$ 训练 FLOPs”估算，同样是把主要线性层成本与参数数目联系起来。它适合快速做预算；遇到长上下文注意力项或大输出头，仍要回到具体模型，检查这些成本是否已被正确计入。

7.4 权重字节数只是内存的第一项

计算量有了大致范围，再来看模型能不能放进内存。若有 $N$ 个参数，每个元素 $s$ 字节，权重占用约 $N s$ 。float32 是 4 字节，float16/bfloat16 通常是 2 字节。本书模型 $N = 7264$ ，纯 float32 参数为 29,056 字节。除以 $1024^{2}$ 得到 MiB；除以 $10^{6}$ 得到 MB，不应把两个单位混在一起。

但训练时，内存里不只有权重。反向需要梯度和保存下来的激活，优化器需要状态，执行算子还可能需要临时工作区；分配器也会保留一些空间。以全 float32 AdamW 的简化账本为例，权重、梯度、一阶动量和二阶动量各占一份，共约 $16 N$ 字节，这还没算激活。

混合精度的情况要看具体实现，有的实现会额外保留主权重。因此，半精度权重每个元素只占两字节，并不能保证总训练内存也减半。

显式注意力分数是 [B,H,T,T]，元素数 $B H T^{2}$ 。若 $B = 2, H = 12, T = 1024$ ，一份 float32 分数张量就是 $2 \times 12 \times 1024^{2} \times 4$ 字节，约 96 MiB；训练可能还保存概率或 dropout 相关状态。融合注意力通过改变存储和计算安排，可能减少这类中间张量，但不能在未运行相应后端时报告具体节省量。

7.5 thop 的数字要问“数到了哪些算子”

手算过预算，再看工具的输出就有了参照。上游 profile(model, inputs=...) 返回 MACs 和参数数，再用 2*macs 换算 FLOPs。这类工具通过模块 hook 观察执行过程：对已注册计数规则的 Linear 等模块，通常容易统计；直接写在模型里的矩阵乘法、自定义模块或某些函数式算子，则可能没有对应处理。

所以，拿到一个总数之后，还要确认它究竟覆盖了哪些计算。

可以先挑一个小模块手算，再与工具结果比较。对注意力，尤其要看 $Q K^{⊤}$ 与 $A V$ 的 $T^{2}$ 项有没有计入：如果工具只捕捉到线性层，报告的运算量就会过分接近随 $T$ 线性增长，漏掉长序列带来的额外成本。

上游实验还把模型转成了 bfloat16。每个数占用的字节更少，不会让乘加次数减半；变化的是数据类型，以及硬件处理这种类型时可能达到的吞吐。

7.6 二分搜索最大批量能证明什么

知道单次计算大致需要多少资源后，还可以尝试寻找设备能容纳的批大小。上游的前向搜索从可选区间中取中点，创建输入和模型，再做前向计数。成功就提高下界；遇到 out-of-memory 就降低上界，并释放引用、清理缓存。

这个搜索结果只适用于当时的执行路径和设备状态。换一个训练流程，最大批量也可能改变。

例如，只做前向时能放下的批量，加上 backward、优化器状态、验证和生成后未必还能放下。empty_cache() 释放的是分配器未被张量引用的缓存，不能删除仍活着的计算图。若异常不是 OOM，应继续抛出，不能把 dtype 不支持或形状错误也当内存不足。

在共享或服务环境反复逼近 OOM 还会影响其他任务。本书不执行这条路径，读者应先用保守估算和小批量检查，再在自己控制的设备上测量，并为运行波动留余量。

7.7 吞吐和 MFU 的量纲检查

批量能放下之后，才轮到“跑得有多快”。若一步处理 $B T$ 个 token，耗时 $Δ t$ 秒，吞吐是 $B T / Δ t$ token/s。若每 token 估算需要 $f$ FLOPs，设备针对相同精度和算子类型的峰值为 $P$ FLOP/s，则理想吞吐为 $P / f$ ，模型 FLOPs 利用率为

MFU = \frac{B T / Δ t}{P / f} = \frac{B T f}{Δ t P} .

这个比值能否解释，取决于分子和分母是否采用同一口径。若计时包括前后向，分子却只填前向 FLOPs，MFU 就会被低估。峰值也必须对应实际使用的精度、稀疏或稠密模式和 GPU 型号，否则算出的比例无法用来比较。

上游的 MFU 示例会同步 GPU 后开始计时，执行前向和反向，再次同步后结束计时；它用前向 FLOPs 的三倍估计前后向总量，没有在这段计时中执行优化器更新。get_gpu_model 则用设备名称子串查硬编码表，取得峰值后计算利用率。表格是实验设置，不是当前硬件规格的权威保证；尤其不能因为字典里给某种 dtype 填了一个值，就推断对应硬件支持它。正式报告应查厂商对应型号、精度、稀疏/稠密条件的规格，并记录计数器覆盖。本书没有据此发布 GPU MFU 数字。

7.8 用预算决定下一步，而不代替测量

现在回看前面的公式，就能解释一些性能变化：当 $T$ 较小、 $D$ 较大时，线性层的 $T D^{2}$ 项可能占主导；随着 $T$ 增长，注意力的 $T^{2} D$ 项和中间存储会越来越重要。生成时，如果能够复用旧位置的计算、每次只处理新增位置，成本结构又会改变。下一章的 KV cache 就是沿着这个思路展开的。

预算最有用的用途是排除明显不可能的方案、解释性能随尺寸变化的方向，并设计测量范围。它不能替代真实吞吐、延迟分布和峰值内存。单用户首 token 延迟、逐 token 延迟和多请求总吞吐也应分别报告。

7.9 练习与答案

补充题一：保持其他量不变， $T$ 翻倍时，线性投影和显式注意力分数分别增长多少？答案：前者两倍，后者四倍。

补充题二：一个工具报告 10 GMAC，一次前向用了 0.1 秒。按本章约定，前向运算率是多少？答案：20 GFLOPs / 0.1 秒 = 200 GFLOP/s；不是 200 GFLOPs。

补充题三：纯权重占 2 GiB，为什么 4 GiB 显存可能仍无法训练？答案：梯度、优化器状态、激活、临时空间和框架开销还未计入。必须给完整执行路径做账。

参考与覆盖

FLOPs notebook：四档配置、thop、二分批量、异常清理、设备名称识别、同步计时、每 token FLOPs 与 MFU。
依赖清单：性能工具依赖，未在本书共享环境新增安装。
参数数量可与第 6 章及现存 qa/core_training/results.json 中的模型记录核对；本章 FLOPs 来自显式近似推导。本轮只核对源码、记录与算术，没有重新运行模型测试，也没有运行 GPU 性能测试。

第 07 章 参数、FLOPs、显存与吞吐怎样估算 ​

7.1 从一个矩阵乘法开始记账 ​

7.2 一个普通 Transformer block 的主要运算 ​

7.3 为何反向常粗估为前向的两倍 ​

7.4 权重字节数只是内存的第一项 ​

7.5 thop 的数字要问“数到了哪些算子” ​

7.6 二分搜索最大批量能证明什么 ​

7.7 吞吐和 MFU 的量纲检查 ​

7.8 用预算决定下一步，而不代替测量 ​

7.9 练习与答案 ​

参考与覆盖 ​