像读流水线一样读Transformer Block

第四部分结束在一个阶段。解码器是一叠 Block，GPT-2 small 有十二个，每个 Block 中注意力是两个阶段之一。本文用读流水线的方式读懂整个 Block：什么流进去、每个阶段写什么、参数和成本落在哪里。

它们不会落在同一个地方。

我在 nanoGPT 的 model.py 中读取 Block，空白处标注了形状，然后把 GPT-2 的权重加载进去。接着我故意把它弄坏：逐个拆掉 Block，再让每个 Block 覆盖自己的输入而不是加上去。

这是面向转向 ML 工程的工程师的系列文章的一部分。每一部分都指出一个不像软件那样运作的东西。这一部分是成本在哪里：不在代码里，在别处。

Block 是四行代码

nanoGPT 是两个文件里的最小化 GPT-2。它的 README 现在指向一个后继者，但 model.py 仍然是最简明的 GPT-2 Block 解读。Block.forward 在 model.py:103-106：

def forward(self, x):                   # x: [B, T, 768], the stream
    x = x + self.attn(self.ln_1(x))     # read a normalised copy, add attention's output
    x = x + self.mlp(self.ln_2(x))      # read again, add the feed-forward's output
    return x                            # [B, T, 768]: the shape it came in with

在这两个调用内部，model.py:52-92，有六个张量值得命名。下面的代码是精简过的，没有缩放、掩码或 Dropout。空白处的形状来自用 GPT-2 small 的形状运行的模型，d = 768，12 个头：

qkv = self.c_attn(x)                         # [B, T, 2304]    q, k, v side by side
q, k, v = qkv.split(768, dim=2)              # three views, [B, T, 768] each
q = q.view(B, T, 12, 64).transpose(1, 2)     # [B, 12, T, 64]  12 heads of 64; k, v alike
att = (q @ k.transpose(-2, -1)).softmax(-1)  # [B, 12, T, T]   Part 4's square
y = (att @ v).transpose(1, 2).reshape(B, T, 768)
y = self.c_proj(y)                           # [B, T, 768]     attention out
h = self.gelu(self.c_fc(x))                  # [B, T, 3072]    the feed-forward's hidden layer
out = self.c_proj(h)                         # [B, T, 768]     feed-forward out

加上流本身，那就是这个 Block：x、投影、分数、注意力的输出、隐藏层，以及前馈的输出。其中三个共享流本身的形状。其他都是视图、逐元素步骤，或者是每个头的 [B, 12, T, 64] 布局。

有一种分工值得读两遍。注意力是唯一混合位置的阶段：其输出的第 t 行由流的第 0..t 行构建。LayerNorm 和前馈一次只处理一个位置。

LayerNorm 减去每个 768 数字行的均值并除以标准差，然后应用学习的缩放和偏移；model.py:27 是一次 F.layer_norm 调用。前馈将每行扩展到 3,072，应用 GELU（一个平滑版的 max(0, x)），再缩回来。如果你写过对行做 map，你就写过前馈的控制流。

原文有一个小部件存活不到粘贴进信息流：在那里选前馈隐藏层并读取其参数，然后选分数，那个没有参数。拖动 T，参数不动而算术动。这个分割是最后两节的核心。

加号就是架构

再读这两行。没有一个阶段替换 x。每个阶段都读取一个归一化副本，把自己的计算结果加到原始值上。进入一个 Block 的流在离开时会有两处编辑，下一个 Block 读取结果。

这个从 embedding 到最后一个 Block 的流 [B, T, 768]，文献称之为残差流；本文叫它总线。每个阶段都读取它并写入它；没有阶段拥有它。如果你写过 total += delta 循环，你就知道这个契约：每个阶段都加到运行中的总数上，没有谁能撤回前面阶段加进去的东西。

数字表明契约在生效。我用一篇早期文章的 1,024 个 token 运行 GPT-2 small，测量了长度——一个向量的 norm——取位置上的中位数。在第 1 到 10 个 Block 中，每个阶段添加的向量长度是其读取的流的 16% 到 40%。输出的流指向的方向几乎和输入一样：余弦相似度，1 表示相同方向，在 0.93 到 0.97 之间。norm 也随着流动增长，从 Block 0 后的 49 到最后的 480。总线上没有任何东西被原地归一化。每个阶段归一化自己的副本，ln_f 在输出层把最后一行变成 logits（下一个 token 的分数）之前再归一化一次。

那是一个决定。称 GPT-2 的顺序为 pre-LN：在分支内部归一化，x + attn(ln(x))，不动总线。原始 transformer 和 GPT-1 是 post-LN：在加法之后归一化，x = ln(x + attn(x))，这在每个阶段都对总线进行了重新缩放。GPT-2 的论文用一句话记录了这个改动——同一句也加上了 ln_f——引用了一个先例但没有给出自己的理由。

一个小测试展示了其中一个理由。我用这个博客的文本训练了两个六 Block 模型，一次一个字符，同样的设置没有 warmup（warmup：从小的步长开始然后逐步增加）。Loss 是模型对每个下一个 token 的平均惊讶；越低越好，0 意味着确定且正确。

300 步后 pre-LN 模型达到了约 2.4，好于仅知道前一个字母得到的 2.56。post-LN 模型卡在 3.2 大致是字母频率单独给出的分数，一个种子跑到 600 步仍在原地。给定 100 步 warmup，它在 200 到 250 步之间突破了。这是一个玩具，只有两个种子和从训练批次读取的 loss，但方向与已发表的分析一致——post-LN 为什么需要 warmup。

回到 GPT-2 small：我删除了 Block。完整模型在这 1,024 个 token 上的 loss 是 3.60。删除第 1 到 10 中任意一个，让流直接穿过，loss 在 3.61 到 3.89 之间。删除 Block 0 则跃升到 9.15。

覆盖比删除更糟。保留一个 Block，但让它的输出替换流而不是加上去。对于第 1 到 8 个 Block，loss 高于 9.6。对于第 2 到 7 个 Block 高于 10.8，即词汇表上均匀猜测的分数。比猜测还差。

原文的第二个小部件画出了这个：切换到覆盖模式，第 2 到 7 个 Block 跨过标记均匀猜测的虚线。然后选 Block 0，它几乎不动。

Block 0 是例外，norm 说明了原因。Embedding 进入时 norm 为 4.6，Block 0 的注意力写入一个长度七倍的向量。Block 0 之后，流基本上是 Block 0 的输出，所以它更像是第二次 embedding 而不是一次编辑。在那里覆盖没有代价，3.59，因为加法已经相当于替换了。

三分之二的权重

你可以在一个 Block 里手工数参数，d = 768，feed-forward 宽度 f = 4d。注意力的两个投影是 4d² + 4d：2,362,368 个数字。前馈的两个是 2·d·f + f + d，即 8d² + 5d：4,722,432。两个 LayerNorm 是 3,072。在每个 GPT-2 大小中，前馈占每个 Block 的三分之二，因为比例来自 model.py:82 的 4 * config.n_embd，而不是宽度。Llama 3.1 8B 改变了两者，前馈占一个 Block 的 81%。

十二个 Block 是 8,500 万参数，GPT-2 small 有 124,439,808。其余几乎全是 token 表：3,860 万个数字，占模型的 31%。它被用了两次：第二次是 Part 2 底部的查找，顶部的投影到 logits；model.py:138 把它们绑在一起，同一个数组而不是一份拷贝。在 GPT-2 XL 中同一个表只占模型的 5.2%。小模型把近三分之一的权重花在词汇表上；大模型几乎注意不到它。

成本在别处

在软件中，函数的大小和运行时间来自同一份代码。这里它们分开了。参数主要落在前馈。算术跟着 T 走。

一次前向传播每个权重矩阵中的每个数字大约花费两次操作，一次乘一次加；Part 4 把它们计为 FLOPs。称这些数字的计数为 N，所以参数每个 token 花费 2N。加上注意力项，4·L·T·d（L 是 Block 数），用于分数和整个方阵上的加权和。在 GPT-2 small 的 1,024 token 时，这相当于每个 token 来自参数的 2.47 亿次操作和来自注意力的 3,800 万次。torch 的 FlopCounterMode 在注意力写出时同意这个数字，Part 4 的 eager 路径：每个 token 284.8 百万，对公式的 285.1 百万。

在 sdpa（默认）上，同一个计数器在 256 和 1,024 token 都报告了 2.47 亿。它没有融合 CPU 注意力内核的条目，即 sdpa 调用的单一例程而不是四行代码。注意力项永远不会达到它的总数，但没有任何东西告诉你这一点。

在一个 Block 内部，在 1,024 token 时注意力项占算术的 18%，在 6,048 token（6 倍 d）时超过线性层。时钟也证明了这一点。在我的笔记本 CPU 上，一个 Block 的注意力阶段，包含投影，在 256 token 时占 Block 时间的 38%，1,024 token 时 47%，4,096 token 时 57%，而它只持有三分之一的参数。

公平地看待这个分割。在 GPT-2 构建时的长度上，线性层做了大部分算术，而前馈是其中较大的那一半。注意力项只有在超过 GPT-2 从不接受的长度时才接管。然而当你挑选一个模型时你挑的是两个预算。参数决定内存，加载时就固定；上下文决定成本和越来越大的内存份额，而且每次请求都会带来。

ML 工程师看什么

四个习惯，前两个是算术。

按组件数参数，而不是总计。每个 Block：注意力 4d²，前馈 2·d·f。加上一次词汇表乘 d，用于表。如果你比较一个 124M 模型和一个 1.5B 模型，问问每个里面有多少是词汇表。

估算成本为每个 token 2N，训练为 6N。训练是一次前向和一次反向传播，而反向（计算梯度）大约花费两个前向。这就是每个 token 6N，跑 D 个 token 的一次运行 6·N·D。nanoGPT 的 GPT-2 配置 config/train_gpt2.py 每次迭代跑 491,520 个 token，跑 600,000 次迭代：2,950 亿个 token。6·N·D 是 2.2 × 10²⁰ 次操作，在 A100 的 bf16 峰值时需要 195 小时，一个八卡节点一天。文件说预期约五天，README 说四天，这说明运行时间在峰值的五分之一到四分之一。nanoGPT 自己的估算器加上了注意力项，model.py:296，在 1,024 token 时该项再多 15%。

深度是延迟。同样的 8,500 万 Block 参数有四种花法：6 个宽度 1,088 的 Block，12 个 768，24 个 544，48 个 384。在 batch 1 时，一个 token 穿过堆栈分别花费 3.2、4.0、4.7、6.7 毫秒（我的 CPU）。每个 Block 增加一个固定开销，Block 依次运行。如果你一次服务一个请求，深度就是你等待的时间，即使算术相等。

测量一个 Block 再丢弃它。Layer pruning（发版时移除一些 Block）和 early exit（在答案看起来稳定时停在中间某个 Block）都依赖总线，而总线是不均匀的。移除 GPT-2 small 的一个中间 Block 会使 loss 增加 0.5% 到 8%；移除 Block 0 增加 155%，最后一个 Block 增加 41%。上面的数字在笔记本上跑了一分钟，用的是这个博客的文本。在你自己的流量上跑出来。

这些都不是建模工作。Block 是围绕一个和的四行代码，而这个和让十二个堆叠像一个流水线一样运作。权重在哪里和时间在哪里是两个不同的答案，而且都是加载任何东西之前就能用算术做出来的。

十二部分中的第五部分。一个 Block 通过两个阶段——注意力和前馈——读取流，并把两个结果加回去：加号就是为什么移除一个中间 Block 只花费几个百分点的 loss 而不是整个模型。每个 Block 的权重有三分之二在前馈里；注意力项的算术份额随上下文增长。下一部分读 GPT-2 在堆栈周围做的五个决定，并在该部分出的时候落在系列索引上。如果你直接到这里，Part 1 读分词器，Part 2 读 token id 索引的表，Part 3 读形状，Part 4 读这个 Block 围绕的阶段。

Written with AI assistance. See AI policy.