详解 Transformer 前向传播和训练步骤的 FLOPs 计算:推理约 2N×T,训练约 6N×T(反向传播约两倍工作量),忽略 LayerNorm 等元素级操作。
稠密 Transformer 的一次前向传播,每个参数、每个 token 大约需要两次浮点运算。这一条规则展开来,能告诉你一个回答在算术层面的成本,而它的失效模式几乎能告诉你关于推理的所有有趣的事情。
Transformer 中几乎全部算术运算都是矩阵乘法,几乎所有权重都分布在这些矩阵中。输入向量与权重矩阵相乘时,每个权重恰好被访问一次,每次访问包含一次乘法和一次加法——即一次乘加运算,按照惯例计为两次浮点运算。
Forward pass (inference): FLOPs ~= 2 * N * T
Training step: FLOPs ~= 6 * N * T
N = number of parameters
T = number of tokens processed
The 6 for training is 2 for the forward pass plus roughly 4 for the
backward pass, which computes gradients with respect to both the inputs
and the weights and therefore does about twice the work of the forward.
两种形式都是近似值,忽略了 Layer Normalization、激活函数、Softmax 和残差连接。这些都是对激活值的逐元素操作,而非对权重的操作,在任何规模的模型中,与矩阵乘法相比都可以忽略不计。但这条规则不会无害地忽略注意力机制,下文会专门讨论。
以一个 700 亿参数的稠密模型、1000 个 token 的提示词和 500 个 token 的回答为例。这 1500 个 token 中的每一个都要完整地穿越整个权重堆栈一次。
Model: N = 70e9 parameters, dense
Request: 1,000 input tokens, 500 output tokens
Prefill (the prompt, processed in one parallel pass):
2 * 70e9 * 1000 = 1.40e14 FLOPs = 140 TFLOP
Decode (500 sequential passes, one token each):
2 * 70e9 * 500 = 0.70e14 FLOPs = 70 TFLOP
Total per answer:
2 * 70e9 * 1500 = 2.10e14 FLOPs = 210 TFLOP
一个短回答需要 210 TFLOP。作为参照,这大约是一块高端消费级显卡在理论峰值下几秒钟内能完成的算术量,这也立即解释了为什么大规模推理是硬件问题而非软件问题。
注意,输入 token 和输出 token 的 FLOPs 消耗是相同的。它们消耗的时间和费用并不相同,原因见下一节关于延迟的讨论——这也是提供商为何对输入和输出分别定价的原因。
注意力机制不是与权重相乘,而是查询与键相乘,两者都是随序列长度增长的激活值。因此,每个生成 token 的成本与上下文中已存在的 token 数量成比例,标准估算会加上第二项:
FLOPs per token ~= 2 * N + 2 * n_layers * n_ctx * d_model
n_layers = transformer blocks
n_ctx = tokens already in the context
d_model = model width (hidden size)
有趣的问题是第二项何时会追上第一项,因为在这个临界点之前可以完全忽略注意力机制,而超过这个临界点就不行了。令两者相等并求解:
2 * n_layers * n_ctx * d_model = 2 * N
n_ctx = N / (n_layers * d_model)
For a 70B model with 80 layers and d_model = 8192:
n_ctx = 70e9 / (80 * 8192)
= 70e9 / 655,360
= 106,811 tokens
Check the two ends:
at 1,000 tokens of context: 2*80*1000*8192 = 1.31e9 vs 2N = 1.40e11
attention is 0.9% of the total
at 128,000 tokens of context: 2*80*128000*8192 = 1.68e11 vs 2N = 1.40e11
attention is now the larger term
因此,对于这种形状的模型,注意力成本在约十万 token 上下文之前可以忽略不计,超过之后则成为主导。这一数字解释了长上下文定价为何呈现如此行为、为何百万 token 的上下文窗口在长度上不成比例地昂贵,以及为何当提示词非常长时,Prefill 成本就不再与提示词长度成线性关系。
上述的层数和宽度是 70B 稠密模型的一个合理形状,用来使代数计算具体化。用你关心的模型卡片上的真实数值代入——需要带走的是交叉点公式,而不是 106,811 这个具体数字。
FLOPs 计数除以一个速率才有意义。加速卡数据手册标注的是峰值,而实际工作负载只能达到其中一部分;这个比例称为 Model FLOPs Utilization,对于大型、调优良好的训练运行,已公布的值大致在 35% 到 50% 之间。推理的 Prefill 阶段可以达到类似水平;推理的 Decode 阶段则不能,原因见下一节。
NVIDIA's published specification for the H100 SXM gives a dense BF16
peak of about 989 TFLOP/s (the doubled figure some tables show assumes
structured sparsity, which general inference does not use).
Take 40% utilisation as a working assumption:
effective rate = 0.40 * 989e12 = 3.96e14 FLOP/s
time for our 210 TFLOP answer:
2.10e14 / 3.96e14 = 0.53 seconds of one accelerator's arithmetic
一个 500 token 的回答需要半秒的 GPU 时间。这是成本端的计算,也是服务经济学的基石:提供商在加速卡上填充大量并发请求,使这半秒的工作能够同时为数十个用户服务。
上面的半秒是算术时间,而等待这个回答的单个用户实际等待的时间要长得多。生成一个 token 需要将模型的权重从内存中读取出来,而在 batch size 为 1 时,加速卡几乎全部时间都在等待内存而非计算。FLOPs 计数不是绑定约束;内存带宽才是。
Same 70B model at BF16, batch size 1, on the same accelerator.
Bytes that must be read per generated token:
70e9 params * 2 bytes = 1.40e11 bytes
At the H100 SXM's published 3.35 TB/s of memory bandwidth:
1.40e11 / 3.35e12 = 0.0418 s per token -> ~24 tokens/second
500 output tokens -> ~21 seconds of wall clock,
against 0.53 seconds of arithmetic.
The arithmetic units are idle roughly 97% of the time.
这是 FLOPs 计数最不能告诉你的事情,也是批处理存在的原因:添加第二个并发请求几乎不增加额外的内存流量,因为同一组权重已经在被读取了。吞吐量随 batch size 扩展,直到算术运算终于成为约束。完整版本的计算见 tokens-per-second 页面。
混合专家模型。使用活跃参数而非总数。一个总参数 400B、活跃参数 40B 的模型,每个 token 执行 2 * 40e9 * T FLOPs,但需要为全部 400B 加载内存。将规则应用于标题数字会将算术量高估十倍——两个参数数量并存的原因恰恰在于此。
推理模型。规则计算 token 数,而推理模型会发出你看不到的 token。显示 200 token 的回答可能实际花费了 4000 token,因此可见 token 的 FLOPs 可能差一个数量级。将推理 token 也计为输出。
缓存前缀。缓存命中会完全跳过缓存跨度的 Prefill 算术运算。FLOPs 没有减少,只是不再重复执行。
量化不改变 FLOPs 计数。以 INT4 而非 BF16 运行,在更廉价的运算单元上执行相同数量的运算。它改变的是移动的字节数和可用的每秒运算数,而非所需的运算数。
投机解码完全改变了计量方式。草稿模型提出若干个 token,大模型在一次前向传播中验证它们,因此产生的 token 数和完整前向传播次数不再是同一数量。
在这此限制内使用,2 * N * T 是应用机器学习中最有用的单一方程式。它为硬件规划提供规模、约束成本、解释定价,而且与这个主题上几乎所有其他数字不同,它不会过时。
How Fast Models Generate: The Roofline and What It Cannot Tell You
How Big a Model Is on Disk: Parameters Times Bytes Per Weight