专栏
/ Scaling Book 中文版 / 第 4 章 第 4 章 6762 字 · 16 min read
查看本章英文源码 ↗
数点
先从具有以下形状的向量 x x x y y y A A A B B B
array shape x [P] y [P] A [N P] B [P M] \def \red#1{\textcolor{red}{#1}}
\def \green#1{\textcolor{green}{#1}}
\def \blue#1{\textcolor{blue}{#1}}
\def \purple#1{\textcolor{purple}{#1}}
\def \orange#1{\textcolor{orange}{#1}}
\def \gray#1{\textcolor{gray}{#1}}
\begin{array}{cc}
\textrm{array} & \textrm{shape} \\ \hline
x & \textrm{[P]} \\
y & \textrm{[P]} \\
A & \textrm{[N P]} \\
B & \textrm{[P M]} \\
\hline
\end{array} array x y A B shape [P] [P] [N P] [P M]
点积 x ⋅ y x \cdot y x ⋅ y P P P 加法 和乘法 ，总计 2 P 2P 2 P
矩阵—向量乘积 A x Ax A x N N N A A A 2 N P 2NP 2 NP
矩阵—矩阵乘积 A B AB A B M M M B B B 2 N P M 2NPM 2 NPM
一般而言，如果有两个高维数组 C C C D D D C [ G H I J K L ] , D [ G H M N K L ] C[\blue{GH}IJ\red{KL}], D[\blue{GH}MN\red{KL}] C [ G H I J K L ] , D [ G H MN K L ] C C C D D D 2 G H I J M N K L 2\blue{GH}IJMN\red{KL} 2 G H I J MN K L 1
Operation FLOPs Data x ⋅ y 2 P 2 P A x 2 N P N P + P A B 2 N P M N P + P M [ c 0 , . . . , c N ] ⋅ [ d 0 , . . . , d N ] 2 ∏ c i × ∏ d j ∉ B A T C H d j ∉ C O N T R A C T d j ∏ c i + ∏ d j \begin{array}{ccc}
\textrm{Operation} & \textrm{FLOPs} & \textrm{Data} \\
\hline
x \cdot y & 2P & 2P \\
A x & 2NP & NP + P \\
AB & 2NPM & NP + PM \\
[c_0,...,c_N] \cdot [d_0,...,d_N] &
2 \prod c_i \times \prod_{\substack{d_j \notin \blue{BATCH} \\ d_j \notin \red{CONTRACT}}} d_j
&
\prod c_i + \prod d_j \\
\hline
\end{array} Operation x ⋅ y A x A B [ c 0 , ... , c N ] ⋅ [ d 0 , ... , d N ] FLOPs 2 P 2 NP 2 NPM 2 ∏ c i × ∏ d j ∈ / B A TC H d j ∈ / CONTR A CT d j Data 2 P NP + P NP + PM ∏ c i + ∏ d j 请记住：对于矩阵—矩阵乘法，计算量 按三次方 O ( N 3 ) O(N^3) O ( N 3 ) O ( N 2 ) O(N^2) O ( N 2 ) 更容易 达到计算饱和上限。这种情况极其罕见，也在很大程度上解释了为什么我们会采用以矩阵乘法为主的架构——它们很适合扩展！
前向与反向 FLOP
训练时，我们并不特别关心某次矩阵乘法的结果；真正关心的是它的导数。事实证明，计算这个导数的成本大约是只执行矩阵乘法本身的 3 倍。
假设 B 只是更大网络中的一个矩阵，A 是输入激活值，并且 C = A B ，那么根据链式法则，损失 L 对 B 的导数为：
∂ L ∂ B = ∂ L ∂ C ∂ C ∂ B = A T ( ∂ L ∂ C ) \frac{\partial L}{\partial B} = \frac{\partial L}{\partial C}\frac{\partial C}{\partial B} = A^T \left(\frac{\partial L}{\partial C}\right) ∂ B ∂ L = ∂ C ∂ L ∂ B ∂ C = A T ( ∂ C ∂ L ) 计算它需要 2 N P M 2NPM 2 NPM N N N A 的导数为
∂ L ∂ A = ∂ L ∂ C ∂ C ∂ A = ( ∂ L ∂ C ) B T \frac{\partial L}{\partial A} = \frac{\partial L}{\partial C}\frac{\partial C}{\partial A} = \left(\frac{\partial L}{\partial C}\right) B^T ∂ A ∂ L = ∂ C ∂ L ∂ A ∂ C = ( ∂ C ∂ L ) B T 这同样需要 2 N P M 2NPM 2 NPM dL/dC 是一个大小为 [ N , M ] [N, M] [ N , M ]
把这些成本相加，可以看到，训练期间总共有 6NPM FLOP ，而推理期间是 2NPM：前向传播为 2NPM，反向传播为 4NPM。因为 PM 是矩阵中的参数数量，所以这就是著名的 Transformer 训练 FLOP 近似式 6 ∗ num parameters ∗ num tokens 6 * \text{num parameters} * \text{num tokens} 6 ∗ num parameters ∗ num tokens 6 ∗ num parameters 6 * \text{num parameters} 6 ∗ num parameters
Transformer 就是未来。好吧，至少它们是现在。也许几年前，它们只是众多架构中的一种。但在今天，几乎值得了解这种架构的每一个细节。这里不会重新介绍该架构，不过可以参考这篇博客 和最初的 Transformer 论文 。
下面是 Transformer 解码器架构的基本示意图：
注意［门控 einsum］： 上图使用了“门控 einsum ”[glu] ，其中把上投影矩阵拆成两个矩阵（上图中的 W In1 W_\text{In1} W In1 W In2 W_\text{In2} W In2 W In W_\text{In} W In
注意 2［MHA 注意力］： 对自注意力而言，T 与 S 相同；但对交叉注意力而言，两者可能不同。在普通多头注意力（Multi-Head Attention，MHA）中，N 与 K 相同；而对于多查询注意力 （Multi-Query Attention，MQA）[mqa] ，K=1；对于分组 MQA （Grouped MQA，GMQA）[gmqa] ，只要求 K 能整除 N。
注意 3［前置归一化与后置归一化］： 上图展示的是所谓“前置归一化”（pre-norm）架构，其中归一化发生在残差连接之前，通常写作 x + attn(norm(x))。如今，LLaMA-3 等模型采用这种架构。最初的 Transformer 论文使用的是“后置归一化”（post-norm）架构，其中 LayerNorm 发生在残差连接之后，即 norm(x + attn(x))。
全局 FLOP 与参数计算
我们来计算 Transformer 每层的 FLOP（这样就不用到处附加 L 因子了）。请注意，下面的训练 FLOP 几乎总是推理 FLOP 的 3 倍，因此把任意总数除以 3，就能得到仅前向传播的成本。
MLP
Transformer 的 MLP 通常由两个输入矩阵乘法和一个输出矩阵乘法组成，两个输入矩阵乘法的结果会逐元素组合：
operation train FLOPs params A [ B , T , D ] ⋅ W i n 1 [ D , F ] 6 B T D F D F A [ B , T , D ] ⋅ W i n 2 [ D , F ] 6 B T D F D F σ ( A i n 1 ) [ B , T , F ] ∗ A i n 2 [ B , T , F ] O ( B T F ) A [ B , T , F ] ⋅ W o u t [ F , D ] 6 B T D F D F ≈ 18 B T D F 3 D F \begin{array}{ccc}
\textrm{operation} & \textrm{train FLOPs} & \textrm{params} \\
\hline \\
A[B,T,\red{D}] \cdot W_{in1}[\red{D}, F] & 6BTDF & DF \\[10pt]
A[B,T,\red{D}] \cdot W_{in2}[\red{D}, F] & 6BTDF & DF \\[10pt]
\sigma\left(A_{in1}\right)[B,T, F] * A_{in2}[B,T, F] & \gray{O(BTF)} \\[10pt]
A[B,T,\red{F}] \cdot W_{out}[\red{F}, D] & 6BTDF & DF \\[10pt]
\hline \\
& \approx 18BTDF & 3DF
\end{array} operation A [ B , T , D ] ⋅ W in 1 [ D , F ] A [ B , T , D ] ⋅ W in 2 [ D , F ] σ ( A in 1 ) [ B , T , F ] ∗ A in 2 [ B , T , F ] A [ B , T , F ] ⋅ W o u t [ F , D ] train FLOPs 6 BT D F 6 BT D F O ( BTF ) 6 BT D F ≈ 18 BT D F params D F D F D F 3 D F
注意力
对于 Q 头数与 KV 头数不同的一般分组查询注意力情形，假设 Q 、K 、V 投影的头维度 H 相同，并估算 QKVO 矩阵乘法的成本：
operation train FLOPs params A [ B , T , D ] ⋅ W Q [ D , N , H ] 6 B T D N H D N H A [ B , T , D ] ⋅ W K [ D , K , H ] 6 B T D K H D K H A [ B , T , D ] ⋅ W V [ D , K , H ] 6 B T D K H D K H A [ B , T , N , H ] ⋅ W O [ N , H , D ] 6 B T D N H D N H 12 B T D ( N + K ) H 2 D ( N + K ) H \begin{array}{ccc}
\textrm{operation} & \textrm{train FLOPs} & \textrm{params} \\
\hline \\
A[B,T,\red{D}] \cdot W_{Q}[\red{D}, N, H] & 6BTDNH & DNH \\[10pt]
A[B,T,\red{D}] \cdot W_{K}[\red{D}, K, H] & 6BTDKH & DKH \\[10pt]
A[B,T,\red{D}] \cdot W_{V}[\red{D}, K, H] & 6BTDKH & DKH \\[10pt]
A[B,T,\red{N}, \red{H}] \cdot W_{O}[\red{N}, \red{H}, D] & 6BTDNH & DNH \\[10pt]
\hline \\ & 12BTD(N+K)H & 2D(N+K)H
\end{array} operation A [ B , T , D ] ⋅ W Q [ D , N , H ] A [ B , T , D ] ⋅ W K [ D , K , H ] A [ B , T , D ] ⋅ W V [ D , K , H ] A [ B , T , N , H ] ⋅ W O [ N , H , D ] train FLOPs 6 BT D N H 6 BT DKH 6 BT DKH 6 BT D N H 12 BT D ( N + K ) H params D N H DKH DKH D N H 2 D ( N + K ) H 点积注意力运算更加微妙：它实际上是一次 T H ⋅ H S TH \cdot HS T H ⋅ H S B B B K K K T S ⋅ S H TS \cdot SH TS ⋅ S H B B B K K K
operation train FLOPs Q [ B , T , K , G , H ] ⋅ K [ B , S , K , H ] 6 B T S K G H = 6 B T S N H softmax S L [ B , T , S , K , G ] O ( B T S K G ) = O ( B T S N ) S [ B , T , S , K , G ] ⋅ V [ B , S , K , H ] 6 B T S K G H = 6 B T S N H ≈ 12 B T S N H = 12 B T 2 N H \begin{array}{cc}
\textrm{operation} & \textrm{train FLOPs} \\
\hline \\[3pt]
Q[\blue{B}, T, \blue{K}, G, \red{H}] \cdot K[\blue{B}, S, \blue{K}, \red{H}]
& 6BTSKGH = 6BTSNH \\[3pt]
\textrm{softmax}_S \;\; L[B, T, S, K, G] & \gray{O(BTSKG) = O(BTSN)} \\[3pt]
S[\blue{B}, T, \red{S}, \blue{K}, G] \cdot V[\blue{B}, \red{S}, \blue{K}, H]
& 6BTSKGH = 6BTSNH \\[3pt]
\hline \\
& \approx 12BTSNH = 12BT^2NH \\
\end{array} operation Q [ B , T , K , G , H ] ⋅ K [ B , S , K , H ] softmax S L [ B , T , S , K , G ] S [ B , T , S , K , G ] ⋅ V [ B , S , K , H ] train FLOPs 6 BTS K G H = 6 BTSN H O ( BTS K G ) = O ( BTSN ) 6 BTS K G H = 6 BTSN H ≈ 12 BTSN H = 12 B T 2 N H 注意［因果掩码］： 最近的大多数 Transformer 使用因果掩码，而非完整的双向注意力。在这种情况下，点积运算的有效 FLOP 会减半。为了在实践中实现这种缩减，我们需要使用注意力内核，而不能采用朴素的 einsum。
其他运算
Transformer 中还会发生其他几种运算。与之相比，LayerNorm 的成本很低，在一阶成本估算中可以忽略。请注意，每层通常包含两个 LayerNorm（一个在注意力之前，一个在 MLP 之前）。此外，还有最后那个规模巨大的反嵌入矩阵乘法（不过它不是逐层出现的）。
operation train FLOPs params 2 × layernorm D A [ B , T , D ] O ( B T D ) 2 D A [ B , T , D ] ⋅ W u n e m b e d [ D , V ] 6 B T D V D V \begin{array}{ccc}
\textsf{operation} & \textsf{train FLOPs} & \textsf{params} \\
\hline \\
2 \times \textrm{layernorm}_D \;\; A[B,T,\red{D}] & \gray{O\left(BTD\right)} & \gray{2D} \\[10pt]
A[B,T,\red{D}] \cdot W_{unembed}[\red{D}, V] & 6BTDV & DV \\
\end{array} operation 2 × layernorm D A [ B , T , D ] A [ B , T , D ] ⋅ W u n e mb e d [ D , V ] train FLOPs O ( BT D ) 6 BT D V params 2 D D V
如果忽略点积注意力的成本（对较短上下文训练而言，这是合理的），那么所有层的总 FLOP 为
( 18 B T D F + 12 B T D ( N + K ) H ) L = 6 ∗ B T ∗ ( 3 D F + 2 D ( N + K ) H ) L = 6 ∗ num tokens ∗ parameter count \begin{align*}
(18BTDF + 12BTD(N+K)H)L = 6 *BT * (3DF + 2D(N+K)H)L \\ = 6 * \textrm{num tokens} * \textrm{parameter count}
\end{align*} ( 18 BT D F + 12 BT D ( N + K ) H ) L = 6 ∗ BT ∗ ( 3 D F + 2 D ( N + K ) H ) L = 6 ∗ num tokens ∗ parameter count 由此得到一个著名的经验法则，用于估算稠密 Transformer 的 FLOP 数量，其中忽略注意力 FLOP。（反嵌入也是一个简单的矩阵乘法，具有 6 B T D V 6BTDV 6 BT D V D V DV D V
注意力成本随上下文长度变化的占比
如果确实计入上面的点积注意力，并假设 F = 4 D F=4D F = 4 D D = N H D=NH D = N H N = K N=K N = K
attention FLOPs matmul FLOPs = 12 B T 2 N H 18 B T D F + 24 B T D N H = 12 B T 2 D 4 ∗ 18 B T D 2 + 24 B T D 2 = 12 B T 2 D 96 B T D 2 = T 8 D \small{\frac{\textrm{attention FLOPs}}{\textrm{matmul FLOPs}} = \frac{12BT^2NH}{18BTDF + 24BTDNH} = \frac{12BT^2D}{4*18 BTD^2 + 24 BTD^2} = \frac{12BT^2D}{96 BTD^2} = \frac{T}{8D}} matmul FLOPs attention FLOPs = 18 BT D F + 24 BT D N H 12 B T 2 N H = 4 ∗ 18 BT D 2 + 24 BT D 2 12 B T 2 D = 96 BT D 2 12 B T 2 D = 8 D T 结论是，训练期间只有当 T>8D 时，点积注意力 FLOP 才会占据主导。 对于 D ~ 8k，这大约对应 64K 个 token。这也合乎情理，因为它意味着，随着 MLP 大小增加，注意力 FLOP 会变得没那么重要。对于大型模型，注意力的二次方成本实际上并不是延长上下文训练的巨大障碍。然而，对于较小模型，例如 D=4608 的 Gemma-27B，注意力会在序列长度约为 37k 时占据主导。2 附录 A 中简要讨论。
其他数学知识
稀疏性与专家混合
如果不简要讨论专家混合（Mixture of Experts，MoE）模型[moe] ，那将是我们的疏忽。MoE 会把标准 Transformer 中单个稠密 MLP 块替换成一组可动态路由的独立 MLP。近似来看，MoE 就是一个普通稠密模型，只不过每层有 E 个 MLP 块 ，而不是只有一个。每个 token 会激活其中 k k k k ≪ E k \ll E k ≪ E E / k E / k E / k DeepSeek v3 实际上有 k = 8 k=8 k = 8 E = 256 E=256 E = 256 O ( E ) O(E) O ( E ) k k k
来源：Deepgram 的混合专家指南 。
与稠密模型相比，MoE 会引入新的通信，主要是两次 AllToAll（一次在 MoE 块之前，一次在之后），用于把 token 路由到正确的专家，再将其送回原设备。3
梯度检查点
反向传播是一种以计算换内存的算法。反向传播无需 O ( n layers 2 ) O(n_\text{layers}^2) O ( n layers 2 ) 需要 O ( n layers ) O(n_\text{layers}) O ( n layers ) ，用于保存前向传播期间产生的所有中间激活值。尽管这好于二次方计算，但内存成本极其高昂：设模型的 B ∗ T = 4 M B * T=4M B ∗ T = 4 M 2 ∗ 20 ∗ B ∗ T ∗ D ∗ L = 84 T B 2 * 20 * B * T * D * L = 84TB 2 ∗ 20 ∗ B ∗ T ∗ D ∗ L = 84 TB
f ( x ) = exp ( g ( x ) ) f(x) = \exp(g(x)) f ( x ) = exp ( g ( x )) d f d x = exp ( g ( x ) ) ⋅ d g d x \frac{df}{dx} = \exp(g(x)) \cdot \frac{dg}{dx} d x df = exp ( g ( x )) ⋅ d x d g 所以，为了避免重新计算，我们需要保存前向传播中的 g ( x ) g(x) g ( x ) exp ( g ( x ) ) \exp(g(x)) exp ( g ( x ))
按块重计算： 只保存每层的输入。这是我们使用的最激进方法，每层只保存 1 个检查点，意味着在上面的例子中只需保存 4.2TB。代价是，反向传播时基本要重复执行前向传播的全部 FLOP，也就是把 FLOP 从 6 ⋅ num params ⋅ num tokens 6 \cdot \text{num params} \cdot \text{num tokens} 6 ⋅ num params ⋅ num tokens 8 ⋅ num params ⋅ num tokens 8 \cdot \text{num params} \cdot \text{num tokens} 8 ⋅ num params ⋅ num tokens 只保存大型矩阵乘法： 另一项简单策略是只保存大型矩阵乘法的输出。这样便可避免在反向传播时重新计算任何大型矩阵乘法，但仍需重新计算其他激活函数和部分注意力。它会把上面的每层 20 个节点减少到更接近每层 7 个。
这绝不是一份全面清单。使用 JAX 时，这些操作通常由 jax.remat/jax.checkpoint 控制（可在这里 阅读更多内容）。
键—值（KV）缓存
正如第 7 节 将会介绍的，LLM 推理有两个关键部分：预填充和生成。
预填充（Prefill） 处理一个长提示词，并把注意力激活值保存到键—值缓存（Key-Value Cache，KV Cache）中，以供生成阶段使用；具体保存的是注意力块中的键—值投影。生成（Generation） 把多个这样的 KV 缓存批处理到一起，并从每个缓存中采样 token。
因此，每个 KV 缓存实际上都是一个大小为 [ 2 , S , L , K , H ] [2, S, L, K, H] [ 2 , S , L , K , H ] 2 S L K H 2SLKH 2 S L KH K H = N H = D = 8192 KH = NH = D = 8192 KH = N H = D = 8192 2 ⋅ 8192 ⋅ 64 ⋅ 8192 = 8 GiB 2 \cdot 8192 \cdot 64 \cdot 8192 = 8\text{GiB} 2 ⋅ 8192 ⋅ 64 ⋅ 8192 = 8 GiB K ≪ N K \ll N K ≪ N
本节应当掌握哪些要点？
Transformer 的总参数量和 FLOP 很容易计算。假设采用 MHA（批大小为 B、词表大小为 V、序列长度为 T、D=dmodel 、F=dff ），总结如下：
组件 每层参数量 每层训练 FLOP MLP 3DF 18BTDF 注意力 4DNH 24BTDNH + 12BT2 NH 其他 2D BTD 词表 DV（总计，并非每层） 12BTDV
只要序列长度 T < 8 D T < 8D T < 8 D
对于合理的上下文长度，训练期间的总 FLOP 预算可以用 6 ⋅ num_params ⋅ num_tokens 6 \cdot \text{num\_params} \cdot \text{num\_tokens} 6 ⋅ num_params ⋅ num_tokens
推理期间，每个 KV 缓存的大小大约为 2 ⋅ S ⋅ L ⋅ K ⋅ H 2 \cdot S \cdot L \cdot K \cdot H 2 ⋅ S ⋅ L ⋅ K ⋅ H
练习题
问题 1： 一个 D = 4096 D=4096 D = 4096 F = 4 ⋅ D F=4 \cdot D F = 4 ⋅ D V = 32 , 000 V=32,000 V = 32 , 000 L = 64 L=64 L = 64 你可以假设 N ⋅ H = D N\cdot H=D N ⋅ H = D
点击此处查看答案。
总参数量大约为 L ⋅ ( 3 D F + 4 D N H + 2 D ) + 2 D V L \cdot (3DF + 4DNH + 2D) + 2DV L ⋅ ( 3 D F + 4 D N H + 2 D ) + 2 D V 64 ⋅ ( 3 ⋅ 4 e 3 ⋅ 16 e 3 + 4 ⋅ 4 e 3 ⋅ 4 e 3 + 2 ⋅ 4 e 3 ) + 2 ⋅ 4 e 3 ⋅ 32 e 3 = 16 e 9 64 \cdot (3 \cdot 4e3 \cdot 16e3 + 4 \cdot 4e3 \cdot 4e3 + 2 \cdot 4e3) + 2 \cdot 4e3 \cdot 32e3 = 16e9 64 ⋅ ( 3 ⋅ 4 e 3 ⋅ 16 e 3 + 4 ⋅ 4 e 3 ⋅ 4 e 3 + 2 ⋅ 4 e 3 ) + 2 ⋅ 4 e 3 ⋅ 32 e 3 = 16 e 9
一般情况下，注意力参数与总参数量之比为 4 D N H / ( 4 D N H + 3 D F ) = 4 D 2 / ( 4 D 2 + 12 D 2 ) = 1 / 4 4DNH / (4DNH + 3DF) = 4D^2 / (4D^2 + 12D^2) = 1/4 4 D N H / ( 4 D N H + 3 D F ) = 4 D 2 / ( 4 D 2 + 12 D 2 ) = 1/4
对每个 token，KV 缓存为 2 ⋅ L ⋅ N ⋅ H = 2 ⋅ 64 ⋅ 4096 2 \cdot L \cdot N \cdot H = 2 \cdot 64 \cdot 4096 2 ⋅ L ⋅ N ⋅ H = 2 ⋅ 64 ⋅ 4096 512 KiB / token。
问题 2： 在 {'X': 4, 'Y': 8, 'Z': 4} 上执行 A[BX , DY ] *D W[DY , F]，总共需要多少 FLOP？每块 TPU 执行多少 FLOP？
点击此处查看答案。
该运算的“理论”FLOP 总数为 2 ⋅ B ⋅ D ⋅ F 2 \cdot B \cdot D \cdot F 2 ⋅ B ⋅ D ⋅ F 2 ⋅ B ⋅ D ⋅ F ⋅ Z 2 \cdot B \cdot D \cdot F \cdot Z 2 ⋅ B ⋅ D ⋅ F ⋅ Z 2 ⋅ B ⋅ D ⋅ F / ( X ⋅ Y ) 2 \cdot B \cdot D \cdot F / (X \cdot Y) 2 ⋅ B ⋅ D ⋅ F / ( X ⋅ Y )
问题 3： 执行 A [ I , J , K , L ] ∗ B [ I , J , M , N , O ] → C [ K , L , M , N , O ] A[I,J,K,L] * B[I,J,M,N,O] \rightarrow C[K,L,M,N,O] A [ I , J , K , L ] ∗ B [ I , J , M , N , O ] → C [ K , L , M , N , O ]
点击此处查看答案。
按照上面的规则，I 和 J 是收缩维度，而 K、L、M、N、O 是非收缩维度。这里没有“批处理维度”，因此结果就是 2 ⋅ I ⋅ J ⋅ K ⋅ L ⋅ M ⋅ N ⋅ O 2 \cdot I \cdot J \cdot K \cdot L \cdot M \cdot N \cdot O 2 ⋅ I ⋅ J ⋅ K ⋅ L ⋅ M ⋅ N ⋅ O
问题 4： 自注意力的算术强度是多少（忽略 Q/K/V/O 投影）？请把答案表示为 Q 和 KV 长度 T 与 S 的函数。 注意力在多长的上下文中会受到 FLOP 限制？给定 TPU 的 HBM 带宽，请绘制注意力相对于 FFW 块的有效成本随上下文长度增长而变化的曲线。
点击此处查看答案。
自注意力需要加载 Q Q Q K K K V V V softmax ( Q ⋅ K ) ⋅ V \text{softmax}(Q \cdot K) \cdot V softmax ( Q ⋅ K ) ⋅ V
Q[B,T,N,H] → reshape Q[B, T, K, G, H] ⋅ K[B, S, K, H] → O[B, T, S, K, G] \text{Q[B,T,N,H]} \rightarrow_\text{reshape} \text{Q[B, T, K, G, H]} \cdot \text{K[B, S, K, H]} \rightarrow \text{O[B, T, S, K, G]} Q[B,T,N,H] → reshape Q[B, T, K, G, H] ⋅ K[B, S, K, H] → O[B, T, S, K, G] U = softmax S ( O[B, T, S, K, G] ) U=\text{softmax}_S(\text{O[B, T, S, K, G]}) U = softmax S ( O[B, T, S, K, G] ) U[B, T, S, K, G] ⋅ V[B, S, K, H] → X[B, T, K, G, H] \text{U[B, T, S, K, G]} \cdot \text{V[B, S, K, H]} \rightarrow \text{X[B, T, K, G, H]} U[B, T, S, K, G] ⋅ V[B, S, K, H] → X[B, T, K, G, H] 因此，总字节数为 2 ∗ sizeof ( Q ) + 2 ∗ sizeof(K or V) = 4 B T N H + 4 B S K H = 4 B H K ∗ ( T G + S ) 2 * \text{sizeof}(Q) + 2 * \text{sizeof(K or V)} = 4BTNH + 4BSKH = 4BHK * (TG + S) 2 ∗ sizeof ( Q ) + 2 ∗ sizeof(K or V) = 4 BTN H + 4 BS KH = 4 B HK ∗ ( TG + S ) 4 B T S N H + O ( B T S N ) 4BTSNH + O(BTSN) 4 BTSN H + O ( BTSN ) 4 B T S K G H / ( 4 B H K ∗ ( T G + S ) ) 4BTSKGH / (4BHK * (TG + S)) 4 BTS K G H / ( 4 B HK ∗ ( TG + S ))
所以基本上，在预填充期间有 S = T S=T S = T 4 B T 2 K G H / 4 B H K T ⋅ ( G + 1 ) = T G / ( G + 1 ) = O ( T ) 4BT^2KGH / 4BHKT \cdot (G+1) = TG/(G + 1) = O(T) 4 B T 2 K G H /4 B HK T ⋅ ( G + 1 ) = TG / ( G + 1 ) = O ( T ) T = 1 T=1 T = 1 4 B S K G H / ( 4 B H K ⋅ ( G + S ) ) = S G / ( G + S ) → G 4BSKGH / (4BHK \cdot (G + S)) = SG / (G + S) \rightarrow G 4 BS K G H / ( 4 B HK ⋅ ( G + S )) = SG / ( G + S ) → G S S S G G G G G G
问题 5： 在什么序列长度下，自注意力 FLOP 会与 QKVO 投影 FLOP 相等？
点击此处查看答案。
这完全是在问 24 B T D N H = 12 B T 2 N H 24BTDNH = 12BT^2NH 24 BT D N H = 12 B T 2 N H 2 D = T 2D = T 2 D = T D = 4096 D=4096 D = 4096 8192 8192 8192
问题 6： 假设在前向传播期间，我们只保存 Transformer 层中 7 次主要矩阵乘法各自的输出（Q、K、V、O + 三个 FFW 矩阵）。为了在反向传播中“重计算”，需要额外执行多少 FLOP？
点击此处查看答案。
只保存七次矩阵乘法的输出（Q、K、V、O、W₁、W₂、W₃），意味着反向传播必须重新计算两次注意力矩阵乘法
Q K ⊤ and softmax ( Q K ⊤ ) V QK^{\top} \quad\text{and}\quad \operatorname{softmax}(QK^{\top})V Q K ⊤ and softmax ( Q K ⊤ ) V 才能得到 ∂ L ∂ W O \frac{\partial L}{\partial W_\text{O}} ∂ W O ∂ L
二者都是 T × T T \times T T × T B B B N N N
4 B T 2 N H . 4 \; B \, T^{2} \, N \, H. 4 B T 2 N H . 其他需要重新计算的运算包括：
执行 O ( B T D ) O(BTD) O ( BT D ) ∂ L ∂ W In1 \frac{\partial L}{\partial W_\text{In1}} ∂ W In1 ∂ L ∂ L ∂ W In2 \frac{\partial L}{\partial W_\text{In2}} ∂ W In2 ∂ L
还要执行 O ( B T F ) O(BTF) O ( BTF ) ∂ L ∂ W Out \frac{\partial L}{\partial W_\text{Out}} ∂ W Out ∂ L
问题 7： DeepSeek v3 称其在 14.8T 个 token 上训练了 2.79M H800 小时（来源 ）。已知它有 37B 个激活参数，它们大约实现了多高的硬件利用率？提示：请注意，它们使用的是不带结构化稀疏的 FP8 FLOP。
点击此处查看答案。
从这里 的规格表可知，带稀疏时的 FP8 性能为 3,026 TFLOP/s；不带稀疏时通常是它的一半（1.513e15 FLOP/s）。2.79M H800 小时意味着总 FLOP 为 2.79e6 * 1.513e15 * 60 * 60 = 1.52e25。给定 37B 个激活参数，这次训练运行应使用约 6 * 37e9 * 14.8e12 = 3.3e24 FLOP。这意味着 FLOP 利用率约为 3.3e24 / 1.52e25 = 21.7%。
问题 8： 专家混合（MoE）模型拥有标准稠密 MLP 块的 E E E k k k k = 8 k=8 k = 8
点击此处查看答案。
因为每个专家有 E E E E ⋅ D ⋅ F E \cdot D \cdot F E ⋅ D ⋅ F k k k 2 ⋅ k ⋅ B ⋅ D ⋅ F 2\cdot k \cdot B \cdot D \cdot F 2 ⋅ k ⋅ B ⋅ D ⋅ F ( 2 ⋅ k ⋅ B D F ) / E D F > 240 (2\cdot k \cdot BDF) / EDF > 240 ( 2 ⋅ k ⋅ B D F ) / E D F > 240 k ⋅ B / E > 120 k \cdot B / E > 120 k ⋅ B / E > 120
因此，要达到计算受限，必须有 B > 120 ⋅ E / k B > 120 \cdot E / k B > 120 ⋅ E / k B > 120 ⋅ 256 / 8 = 3840 B > 120 \cdot 256 / 8 = 3840 B > 120 ⋅ 256/8 = 3840
附录
附录 A：Flash Attention 如何工作？
反对把 Transformer 扩展到极长上下文的传统理由是，注意力 FLOP 和内存使用量会随上下文长度按二次方增长。注意力 QK 乘积的形状确实是 [ B , T , S , N ] [B, T, S, N] [ B , T , S , N ]
正如前文所述，即使成本是二次方的，也只有在 T > 8 ⋅ D T > 8 \cdot D T > 8 ⋅ D
为了计算注意力，我们不需要具体化完整的注意力矩阵！可以计算局部和与局部最大值，从而始终只具体化数组的一小块。尽管 FLOP 总数仍按二次方增长，但内存压力会大幅降低。
第二点最早由 Rabe 等人（2021） 提出，之后又见于 Flash Attention 论文 （Dao 等人，2022）。基本思想是按 K/V 数据块计算注意力：先计算局部 softmax 和一些辅助统计量，再将它们传给下一个数据块，由后者将其与自己的局部数据块合并。具体来说，我们计算
M： 序列维度上 q ⋅ k q \cdot k q ⋅ k O： 序列维度上运行中的完整注意力 softmaxL： 运行中的分母 ∑ i exp ( q ⋅ k i − running max ) \sum_i \exp(q \cdot k_i - \text{running max}) ∑ i exp ( q ⋅ k i − running max )
有了这些量，我们只需恒定大小的内存，就能计算新的最大值、新的运行和以及新的输出。粗略地说，注意力大致执行下面的运算：
Attn ( Q , K , V ) = ∑ i exp ( Q ⋅ K i − max j Q ⋅ K j ) V i ∑ l exp ( Q ⋅ K l − max j Q ⋅ K j ) \text{Attn}(Q, K, V) = \sum_i \frac{\exp(Q \cdot K_i - \max_j Q \cdot K_j) V_i}{\sum_l \exp(Q \cdot K_l - \max_j Q \cdot K_j)} Attn ( Q , K , V ) = i ∑ ∑ l exp ( Q ⋅ K l − max j Q ⋅ K j ) exp ( Q ⋅ K i − max j Q ⋅ K j ) V i 为了数值稳定性，式中减去了最大值；这样做不会影响结果，因为 ∑ i exp ( a i + b ) = exp ( b ) ∑ exp ( a ) \sum_i \exp(a_i + b) = \exp(b) \sum \exp(a) ∑ i exp ( a i + b ) = exp ( b ) ∑ exp ( a ) K 1 K^1 K 1 K 2 K^2 K 2 L 1 L^1 L 1 L 2 L^2 L 2
L 1 = ∑ i exp ( Q ⋅ K i 1 − max j Q ⋅ K j 1 ) L^1 = \sum_i \exp(Q \cdot K_i^1 - \max_j Q \cdot K_j^1) L 1 = i ∑ exp ( Q ⋅ K i 1 − j max Q ⋅ K j 1 ) L 2 = ∑ i exp ( Q ⋅ K i 2 − max j Q ⋅ K j 2 ) L^2 = \sum_i \exp(Q \cdot K_i^2 - \max_j Q \cdot K_j^2) L 2 = i ∑ exp ( Q ⋅ K i 2 − j max Q ⋅ K j 2 ) 那么，利用下面的公式，可以把它们合并成这两个数据块共同的完整 softmax 和：
L combined = exp ( M 1 − max ( M 1 , M 2 ) ) ⋅ L 1 + exp ( M 2 − max ( M 1 , M 2 ) ) ⋅ L 2 L^\text{combined} = \exp(M^1 - \max(M^1, M^2)) \cdot L^1 + \exp(M^2 - \max(M^1, M^2)) \cdot L^2 L combined = exp ( M 1 − max ( M 1 , M 2 )) ⋅ L 1 + exp ( M 2 − max ( M 1 , M 2 )) ⋅ L 2 其中
M 1 = max j Q ⋅ K j 1 and M 2 = max j Q ⋅ K j 2 M^1 = \max_j Q \cdot K_j^1 \text{ and } M^2 = \max_j Q \cdot K_j^2 M 1 = j max Q ⋅ K j 1 and M 2 = j max Q ⋅ K j 2 完整 softmax 也可以采用这种方法，从而让我们能够累积任意大的 softmax 和。下面是 Flash Attention 论文中的完整算法。
从硬件角度看，这让我们能够把 Q 数据块装入 VMEM（也就是上面算法所称的片上 SRAM）；这样每次迭代只需加载 KV 数据块，从而提高算术强度。运行中的统计量也可以保存在 VMEM 中。
最后还要强调一个容易忽略的细节：为了让 Flash VJP（反向模式导数）计算在训练中切实可行，会用到注意力 softmax 的一个性质。我们定义一个中间 softmax 数组：
S i j = e τ q i ⋅ k j ∑ l e τ q i ⋅ k l S_{ij} = \frac{e^{\tau q_i \cdot k_j}}{\sum_l e^{\tau q_i \cdot k_l}} S ij = ∑ l e τ q i ⋅ k l e τ q i ⋅ k j 在注意力中，我们从反向模式的 dO 与 V 数组得到 dS ：
d S i j = d O i d ⋅ d V j d = ∑ d d O i d V j d dS_{ij} = dO_{id} \cdot_d V_{jd} = \sum_d dO_{id} V_{jd} d S ij = d O i d ⋅ d V j d = d ∑ d O i d V j d 在把这个梯度反向传播给 Q 和 K 时，
d ( q i ⋅ k j ) = ( d S i j − S i j ⋅ j d S i j ) S i j d(q_i \cdot k_j) = (dS_{ij} - S_{ij} \cdot_j dS_{ij}) S_{ij} d ( q i ⋅ k j ) = ( d S ij − S ij ⋅ j d S ij ) S ij 我们利用一个恒等式，将沿大型键长度 维度的收缩，换成沿特征深度 维度的局部收缩。
S i j ⋅ j d S i j = ∑ j e τ q i ⋅ k j ∑ k e τ q i ⋅ k k ∑ d d O i d V j d = ∑ d d O i d ∑ j e τ q i ⋅ k j ∑ k e τ q i ⋅ k k V j d = ∑ d d O i d O i d = d O i d ⋅ d O i d \begin{align*}
S_{ij} \cdot_j dS_{ij} &= \sum_j \frac{e^{\tau q_i \cdot k_j}}{\sum_k e^{\tau q_i \cdot k_k}} \sum_d dO_{id} V_{jd} \\
&= \sum_d dO_{id} \sum_j \frac{e^{\tau q_i \cdot k_j}}{\sum_k e^{\tau q_i \cdot k_k}} V_{jd} \\
&= \sum_d dO_{id} O_{id} \\
&= dO_{id} \cdot_d O_{id}
\end{align*} S ij ⋅ j d S ij = j ∑ ∑ k e τ q i ⋅ k k e τ q i ⋅ k j d ∑ d O i d V j d = d ∑ d O i d j ∑ ∑ k e τ q i ⋅ k k e τ q i ⋅ k j V j d = d ∑ d O i d O i d = d O i d ⋅ d O i d 这种替换对实现按序列分块的 VJP 局部 计算至关重要，也进一步支持了环形注意力等巧妙的分片方案。
← 上一章
分片矩阵及其乘法
下一章 →
如何并行化 Transformer 训练
