Transformer 块:LayerNorm、注意力、残差、MLP

Transformer 块:LayerNorm、注意力、残差、MLP

本节摘要:一个块是每个现代 decoder LLM 的单元——LayerNorm、多头注意力、残差、MLP、残差。pre-LN 变体无 warmup 稳定训练,post-LN 是原论文发布的。本节并排构建两者,展示哪个能在常用学习率下扛过 12 层堆叠。每个 decoder-only transformer 块是一个 (B, T, D)→(B, T, D) 的函数,内有两个子层。你会学到两个残差路径与两个归一化位置的选择、因果掩码的必要性、融合 QKV 投影,以及为什么 pre-LN 是 GPT-2 之后的标配。

对应原课程:Phase 19 · Lesson 34 · transformer-block(原英文 phases/19-capstone-projects/34-transformer-block/docs/en.md)。本节属「从零构建 GPT」赛道第五节。

学习目标

阅读完本节,你应当能够:

1. 用四个部件从零建 PyTorch transformer 块:LayerNorm、多头因果注意力、残差连接、位置级 MLP。
2. 把 LayerNorm 放两种配置(pre-LN 与 post-LN),解释为何其一无 warmup 稳定训练。
3. 在多头注意力内实现因果掩码,使 token i 看不到 token j>i。
4. 在 12 层堆叠上跟踪两变体的梯度流,不糊弄地读结果。
5. 把块作 drop-in 单元,供下一节组装 1.24 亿参数 GPT 复用。

一、问题与直觉

transformer 是一个块重复。块错一次,重复 12 次,你就发一个首 epoch 就发散、或余生都需 warmup hack 的模型。本节你会看到两个失败模式(不是异类,首次堆块就冒):注意力层注意未来;LayerNorm 放在深处压不住残差信号的位置。修复是机械的——块正好两条残差路径、两个归一化位置,选对位置,堆叠就只是簿记。

这是 pre-LN 变体——LayerNorm 坐在残差支路内、子层前,残差连接带未归一化信号向前。post-LN 把 LayerNorm 移到残差加之后。形状一致,训练行为不一致:post-LN 下,残差路径回流的梯度得过 LayerNorm,深 12、学习率 3e-4 时,该梯度缩得快到需 warmup;pre-LN 留残差路径未归一化,梯度干净传到嵌入层。pre-LN 正是 GPT-2 之后标配的原因。

二、从零实现

因果多头注意力:子层把输入三路投影成 Q/K/V,各从 (B, T, D) 重塑成 (B, H, T, D/H)。缩放点积注意力算 softmax(QK^T/sqrt(d_k)) 逐头,掩码上三角为负无穷,softmax 应用掩码,乘 V。头拼回单 (B, T, D) 张量再投影一次。掩码是让模型因果的唯一部件——忘掩码就训了个作弊模型。

MLP:位置级 MLP 对每 token 独立施加同两层网络。隐藏宽是嵌入宽的四倍,激活 GELU,第二线性后接 dropout。MLP 内无 token 互通,所有 token 混合发生在注意力里。

残差做两件事:让梯度路径跨深可加(12 层仍保梯度范数在量级),让每块学一个对运行表征的加性更新而非全替换。两效应正是块能 scale 的原因。

三、可复用产物

code/main.py 实现:LayerNorm(可学缩放与偏移、偏 eps、逐 token 向量)、MultiHeadAttention(num_heads、head_dim=d_model//num_heads、融合 QKV 投影、注册因果掩码、注意力与残差 dropout)、FeedForward(两线性层、GELU、dropout)、TransformerBlock(带 pre_ln 标志切两变体)。demo 建 6 层 pre-LN 堆与 6 层 post-LN 堆,同输入,打输出形状与一次反向后嵌入层梯度范数并排——pre-LN 堆的嵌入梯度比 post-LN 堆在同学习率下大一个数量级,这就是 pre-LN 无 warmup 训练的经验信号。

四、生产模式

三个模式把教科书块变成能发的。融合 QKV 投影:三独立线性层三次内核启动三次 matmul;一个宽 3*d_model 的线性层一次启动干同样活,再沿末轴切输出——融合路径在每个加速器上更快,GPT-2/LLaMA/Mistral 的参考实现都这么发。注册因果掩码 buffer:掩码只依赖最大上下文长,构造时 register_buffer 一次,前向切活动窗,省每调用分配;忘了这,长上下文下掩码成分配热点。dropout 两处非三处:dropout 属注意力 softmax 后(注意力 dropout)与 MLP 第二线性后(残差 dropout);残差本身的 dropout 破坏让深度梯度流的加性恒等。早期实现搞错这个,付出脆弱训练的代价。

五、框架对比

pre-LN 是每个现代开权重 LLM 用的,post-LN 是 2017 原注意力论文用的,知道两者够读任何 decoder 架构。把 GELU 换 SiLU 就得 LLaMA 族激活,把 LayerNorm 换 RMSNorm 就得 LLaMA 族归一化——同一骨架。本节块不修改地直插下一节的 GPT 组装。

六、练习

1. 去 bias:给块里每个线性加 bias=False(现代开权重 LLM 不带 bias),量 12 层 768 维模型省多少参数。
2. RMSNorm:把 nn.LayerNorm 换手搓 RMSNorm,确认输出形状不变。
3. 注意力权重检查:加标志返首头 (B, T, T) 权重,画上三角确认 softmax 后为零。
4. 两变体差异:喂 (2, 16, 384)、H=6 过两变体,同初始化、dropout=0,断言前向输出不同(not torch.allclose)。
5. 梯度流:12 层堆,打两变体在嵌入层的梯度范数,确认 pre-LN 大一个数量级。

本节要点回顾

1. 块是单元:transformer = 一个块重复 12 次,块错则全错。
2. 两残差两归一化:pre-LN(子层前归一化,残差带未归一信号)vs post-LN(残差加后归一化)。
3. pre-LN 无 warmup:残差路径未归一化,梯度干净传到嵌入;GPT-2 之后标配。
4. 因果掩码必要:唯一让模型因果的部件,忘掩码训作弊模型。
5. MLP 无 token 互通:所有混合在注意力,MLP 是位置级两层 D→4D→D。
6. 三生产模式:融合 QKV、注册掩码 buffer、dropout 两处非三处。

下一节,我们把 12 个块叠成「GPT 模型组装」——加嵌入、最终 LayerNorm、绑定的 LM 头,得到 1.24 亿参数的 GPT。