1. 从 llm.c 的极简哲学说起为什么第三篇要死磕注意力与 MLP如果你跟着前两篇一路走过来应该已经对 llm.c 这个项目的脾气有所了解了。它不追求工程上的花哨也不搞一堆抽象层而是用最朴素的 C 代码把 GPT-2 从零到一搭出来。到了第三篇我们终于要碰整个 Transformer 架构里最核心、也最容易被讲玄乎的两块骨头多头自注意力机制和MLP 前馈网络。顺带还要把夹在它们中间的LayerNorm彻底讲透。很多人学 GPT 类模型卡就卡在这一步。看论文里的公式觉得懂了看 PyTorch 的nn.MultiheadAttention觉得懂了但一旦让你用 C 手写一遍立刻发现脑子里全是浆糊Q、K、V 到底怎么切分多头到底是多在哪里缩放因子为什么是1/sqrt(d_k)残差连接加在哪个位置LayerNorm 到底对哪个维度做归一化这些问题不解决你永远只是调包侠而不是真正理解 GPT-2 的人。llm.c 的价值恰恰在这里。它逼着你把每一个张量形状、每一次内存布局、每一个循环边界都想清楚。这篇我就按照 llm.c 第三篇的节奏把多头注意力机制原理、MLP 的内部结构、LayerNorm 的归一化维度这三件事从代码视角掰开揉碎讲一遍。适合已经看过前两篇、对 GPT-2 整体结构有概念但一到注意力细节就发怵的读者。如果你还没看过前两篇建议先补一下张量基础和嵌入层那部分不然这篇里的形状推导会让你有点晕。我先把这一篇要解决的核心问题列出来你可以对照自己的理解程度单头注意力到底在算什么为什么需要多头Q、K、V 的投影矩阵在 llm.c 里是怎么存的形状怎么对齐缩放点积注意力的缩放因子从哪来不缩放会怎样因果掩码causal mask在代码里怎么实现为什么不能简单置零MLP 为什么是两层线性加一个 GELU中间维度为什么是 4 倍LayerNorm 的均值方差到底在哪个维度上算和 BatchNorm 差在哪残差连接和 LayerNorm 的先后顺序Pre-LN vs Post-LN对训练的影响这些问题我会一个一个拆。你不需要有深厚的数学背景但需要愿意跟着形状推导走一遍。下面正式开始。2. 多头自注意力机制从单头到多头的形状推演2.1 单头注意力一次加权平均的本质先把最朴素的单头注意力讲清楚。假设输入序列长度为T每个 token 的嵌入维度为CGPT-2 small 里C768。输入张量形状是[T, C]。注意力的目标是让每个位置的 token 去看其他所有位置然后根据相关性做一次加权平均。具体怎么做用三个线性投影把输入分别映射成 Query、Key、ValueQ X W_q形状[T, C]K X W_k形状[T, C]V X W_v形状[T, C]然后计算注意力分数scores Q K^T形状变成[T, T]。这个矩阵的第i行第j列表示第i个 token 对第j个 token 的关注程度。接着做缩放、掩码、softmax得到注意力权重attn形状还是[T, T]。最后out attn V形状回到[T, C]。用一句话概括注意力就是让每个 token 根据自己和其他 token 的相似度对其他 token 的 Value 做加权求和。相似度用 Q 和 K 的点积衡量权重用 softmax 归一化。这里有个细节很多人第一次看会忽略Q K^T得到的是[T, T]而不是[T, C]。也就是说注意力的计算复杂度和序列长度的平方成正比。这就是为什么长上下文那么吃显存——T翻倍注意力矩阵大小翻四倍。llm.c 里这个矩阵是实打实分配出来的你能直观看到它的内存占用。2.2 多头到底多在哪切分维度的选择单头注意力的表达能力有限因为它只学了一种关注模式。多头注意力的思路是把C维的 Q、K、V 切成H份每份维度是C/H每一份独立做一次注意力最后把结果拼回来再投影一次。GPT-2 small 里H12C768所以每个头的维度d_head 768/12 64。这个 64 不是随便定的它要保证H * d_head C这样拼接后维度才能对得上。关键问题来了多头是切分 Q、K、V 的维度还是复制 H 份独立的 Q、K、V答案是前者。在 llm.c 里W_q、W_k、W_v的形状都是[C, C]投影完之后得到[T, C]然后 reshape 成[T, H, d_head]再 transpose 成[H, T, d_head]让每个头独立计算。这样做的原因是如果每个头都用完整的C维做投影参数量会变成H倍而且计算量也爆炸。切分维度是性价比最高的方案。我用一个表格把形状变化列清楚这是理解多头注意力的关键步骤操作形状变化输入X[T, C]QKV 投影X W_qkv[T, 3C]拆分reshape transpose[3, H, T, d_head]注意力分数Q K^T[H, T, T]缩放掩码softmax-[H, T, T]加权求和attn V[H, T, d_head]合并头transpose reshape[T, C]输出投影 W_o[T, C]注意 llm.c 里为了效率通常把W_q、W_k、W_v合并成一个[C, 3C]的大矩阵一次矩阵乘法搞定三次投影。这是个很实用的工程技巧减少 kernel 启动次数也提高内存访问的连续性。你在自己实现的时候可以分开写逻辑更清晰但生产环境建议合并。2.3 缩放因子 1/sqrt(d_head)不是玄学是方差控制1/sqrt(d_head)这个缩放因子很多人背下来了但不知道为什么。我用方差的角度解释一遍。假设 Q 和 K 的每个元素都是独立同分布均值 0方差 1。那么点积q · k sum(q_i * k_i)的方差是多少每个q_i * k_i的方差是1 * 1 1d_head个独立项相加方差就是d_head。所以点积结果的方差是d_head标准差是sqrt(d_head)。当d_head 64时点积结果的标准差是 8。如果不缩放softmax 的输入值域会很大导致 softmax 输出接近 one-hot梯度几乎消失。除以sqrt(d_head)之后方差被拉回 1softmax 的输入分布比较温和梯度更健康。llm.c 里这个缩放通常直接乘在 Q 上或者乘在 scores 上效果一样。我建议乘在 Q 上因为这样只需要对[T, C]做一次缩放而不是对[H, T, T]做省计算。不过 llm.c 为了代码清晰可能是在 scores 上乘的你看具体实现。提示如果你自己写的时候忘了缩放训练初期 loss 会下降得很慢甚至不降因为 softmax 饱和了。这个坑我踩过排查了半天才发现是漏了1/sqrt(d_head)。2.4 因果掩码为什么不能简单置零GPT 是自回归模型预测第t个 token 时只能看到前t个 token不能看到未来。所以在注意力分数矩阵[T, T]里位置(i, j)当j i时必须被屏蔽掉。最常见的做法是把这些位置设成一个很大的负数比如-1e9这样 softmax 之后这些位置的权重接近 0。为什么不能直接设成 0因为 softmax 的公式是exp(x_i) / sum(exp(x_j))如果设成 0exp(0) 1反而会分到权重。必须设成负无穷实践中用大负数近似exp(-1e9) ≈ 0才能真正屏蔽。llm.c 里通常用一个预计算的掩码矩阵或者在上三角区域直接写-INFINITY。我建议用预计算的方式因为掩码是固定的没必要每次前向都重新算。另外要注意掩码是在 softmax 之前加的不是之后。还有个细节softmax 的数值稳定性。如果一行里所有值都是-1e9理论上不会因为对角线ji是允许的或者最大值太大直接exp会溢出。标准做法是先减去每行的最大值再exp。llm.c 里应该做了这个处理你写的时候别忘了。3. MLP 前馈网络4 倍膨胀背后的设计逻辑3.1 两层线性 GELU结构虽简单位置很关键Transformer 每个 block 里注意力后面跟着一个 MLP也叫 FFN前馈网络。GPT-2 的 MLP 结构是Linear(C, 4C) - GELU - Linear(4C, C)。输入[T, C]先升维到[T, 4C]激活后再降回[T, C]。为什么中间要膨胀 4 倍这个问题没有标准答案但有几个合理的解释。第一注意力层主要负责 token 之间的信息交换MLP 负责对每个 token 独立做非线性变换膨胀给模型更大的容量去存储知识。第二4 倍是个经验值实验发现比 2 倍或 8 倍效果好。第三从参数量看MLP 占了整个 block 参数量的三分之二左右2 * C * 4C 8C^2而注意力是4C^2所以 MLP 其实是知识存储的主力。llm.c 里这个 MLP 的实现很直接两个矩阵乘法加一个 GELU。GELU 的公式是x * Phi(x)其中Phi是标准正态分布的累积分布函数。实践中常用近似0.5 * x * (1 tanh(sqrt(2/pi) * (x 0.044715 * x^3)))。llm.c 里应该用的是这个 tanh 近似因为它比精确的 erf 实现快。3.2 GELU vs ReLU为什么 GPT-2 选了 GELUReLU 是max(0, x)简单粗暴负数直接归零。GELU 是平滑的负数区域不是硬归零而是有一个小的负值。这个平滑性带来两个好处一是梯度更平滑训练更稳定二是负值区域保留了一点信息表达能力更强。我用一个生活类比ReLU 像是一个严格的老师答错了直接零分GELU 像是一个温和的老师答错了给个低分但不是零分学生还能从错误中学到东西。当然这个类比不严谨但能帮你记住区别。实测下来GELU 在 Transformer 类模型上确实比 ReLU 好一点点但差距不大。如果你自己搭模型用 ReLU 也能跑只是可能收敛慢一点。llm.c 忠实还原 GPT-2所以用 GELU。3.3 MLP 里的残差和 LayerNorm顺序问题MLP 本身不包含残差和 LayerNorm但它在 block 里的位置和这两个组件紧密相关。GPT-2 用的是 Pre-LN 结构x x Attention(LN(x))然后x x MLP(LN(x))。也就是说LayerNorm 在注意力和 MLP 之前残差连接把输入直接加到输出上。这个顺序很重要。Post-LN原始 Transformer 论文用的是x LN(x Attention(x))LayerNorm 在残差之后。Pre-LN 训练更稳定不需要 warmup 也能收敛所以 GPT-2 之后的大部分模型都用 Pre-LN。llm.c 里你注意看 LayerNorm 的调用位置就能确认是哪种。注意如果你把 Pre-LN 和 Post-LN 搞混训练可能会发散。我建议自己实现时先用 Pre-LN稳。4. LayerNorm 的归一化维度最容易搞错的地方4.1 LayerNorm 到底对哪个维度做归一化这是初学者最容易晕的地方。输入形状[T, C]LayerNorm 是对每个 token 的C维特征做归一化也就是对最后一维做。每个 token 独立计算均值和方差然后归一化再乘gamma加beta。gamma和beta的形状都是[C]对所有 token 共享。对比一下 BatchNormBatchNorm 是对 batch 维度做归一化每个特征维度独立。在 NLP 里序列长度可变batch 里不同样本长度不一样用 BatchNorm 很麻烦。LayerNorm 不依赖 batch每个样本独立所以更适合 NLP。我用一个表格对比维度LayerNormBatchNorm归一化维度特征维 CBatch 维统计量来源单个样本整个 batch是否依赖 batch否是训练/推理一致性一致不一致适合场景NLP、序列CV、固定 batchllm.c 里 LayerNorm 的实现通常是对每个 token遍历C个元素算均值和方差然后归一化。因为C768这个循环不算大但T可能很长所以总体计算量不小。优化的话可以用向量化指令但 llm.c 为了可读性可能用朴素循环。4.2 均值方差的计算数值稳定性不能忽视LayerNorm 的公式是y (x - mean) / sqrt(var eps) * gamma beta。eps是个很小的数通常1e-5防止除零。方差的计算有两种方式一种是E[(x - mean)^2]另一种是E[x^2] - mean^2。后者只需要遍历一次但数值稳定性差当x的均值很大时两个大数相减会损失精度。llm.c 里我建议用前者遍历两次但更稳。还有个细节gamma和beta是可学习参数初始化时gamma1beta0。这样初始状态下 LayerNorm 是恒等变换除了归一化不会破坏输入分布。4.3 LayerNorm 在 llm.c 里的内存布局llm.c 用连续内存存张量[T, C]的布局是行优先第t个 token 的C个特征连续存放。LayerNorm 遍历时外层循环t内层循环c这样内存访问是连续的缓存友好。如果你反过来写性能会差很多。gamma和beta是[C]的一维数组对所有 token 共享。计算时内层循环里同时读x[t*C c]、gamma[c]、beta[c]三个数组都是连续访问效率高。5. 把注意力、MLP、LayerNorm 串起来一个 block 的完整前向5.1 数据流梳理从输入到输出现在把这三块串起来。一个 GPT-2 block 的前向过程是输入x形状[T, C]ln1 LayerNorm(x)形状[T, C]attn MultiHeadAttention(ln1)形状[T, C]x x attn残差连接ln2 LayerNorm(x)形状[T, C]mlp MLP(ln2)形状[T, C]x x mlp残差连接输出x形状[T, C]这个结构重复N层GPT-2 small 是 12 层最后接一个 LayerNorm 和输出投影。注意残差连接的作用它让梯度可以直接回传缓解深层网络的梯度消失。没有残差12 层堆起来训练会很困难。llm.c 里残差就是简单的逐元素加法但要注意原地操作还是新分配内存。原地操作省内存但可能覆盖还需要用的输入要小心。5.2 参数量核算钱花在哪了我算一下 GPT-2 small 一个 block 的参数量让你对钱花在哪有个概念注意力 QKV 投影C * 3C 768 * 2304 1,769,472注意力输出投影C * C 589,824MLP 第一层C * 4C 768 * 3072 2,359,296MLP 第二层4C * C 3072 * 768 2,359,296LayerNorm 两个2 * 2C 3072加起来一个 block 约 708 万参数其中 MLP 占了 471 万约三分之二。12 层就是 8500 万左右加上嵌入层和输出层总共约 1.24 亿参数。这就是 GPT-2 small 的规模。这个核算告诉你如果你想压缩模型MLP 是重点如果你想提升容量加 MLP 宽度比加注意力头数更划算。5.3 常见实现错误清单我在实现和调试过程中踩过的坑列出来帮你避雷形状对不上多头 reshape 时H * d_head必须等于C否则 transpose 后维度错乱。掩码位置错掩码必须在 softmax 之前加且用大负数不是零。缩放漏掉1/sqrt(d_head)忘了乘训练不收敛。LayerNorm 维度错对T维归一化而不是C维结果完全不对。残差覆盖原地加法把还需要用的x覆盖了导致后续计算错误。GELU 近似公式错tanh 里的系数写错激活值偏差大。softmax 溢出没减最大值exp溢出成 inf。这些错误单看都很低级但在 C 代码里没有形状检查一旦写错就是静默错误loss 不降你都不知道哪错了。建议每写一个模块先用小规模数据比如T4, C8, H2手动验证输出。6. 调试与验证怎么确认你的实现是对的6.1 用小规模数据做单元测试llm.c 这种底层实现最怕的就是看起来对但实际错。我的做法是构造极小的输入手动算一遍期望输出然后对比。比如T2, C4, H2, d_head2输入用简单的整数手动算注意力分数、softmax、加权求和然后和代码输出对比。LayerNorm 也可以这样测输入[1, 2, 3, 4]均值 2.5方差 1.25归一化后是[-1.34, -0.45, 0.45, 1.34]近似乘 gamma 加 beta 后对比。6.2 梯度检查数值梯度 vs 解析梯度如果你实现了反向传播梯度检查是必须的。用数值梯度(f(xh) - f(x-h)) / 2h和解析梯度对比误差应该在1e-5量级。llm.c 如果只做前向推理可以跳过这步但如果你想训练这步不能省。6.3 和参考实现对比最靠谱的验证是和一个已知正确的实现对比。你可以用 PyTorch 搭一个同样的 GPT-2 block加载相同的权重输入相同的数据对比每一层的输出。如果每一层都对得上那你的 C 实现就是对的。这个过程中你会定位到具体是哪一层出了问题比盲猜高效得多。提示对比时注意浮点精度C 的 float 和 PyTorch 的 float32 可能有微小差异允许1e-4左右的误差。7. 一些实战心得与性能考量7.1 内存布局决定性能llm.c 用连续内存矩阵乘法的手写实现里循环顺序对性能影响巨大。比如C A BA是[M, K]B是[K, N]C是[M, N]。如果按i, j, k循环内层访问B[k*N j]是跨行访问缓存不友好。改成i, k, j内层访问B[k*N j]连续性能能差好几倍。注意力里的Q K^T也是同理。K^T不是真的转置而是通过索引访问实现循环顺序要配合。7.2 多头注意力的并行化思路llm.c 是单线程 C但多头天然可以并行因为每个头独立。如果你用 OpenMP可以在头维度上并行#pragma omp parallel for加在h循环上。不过要注意每个头写回输出的位置不能冲突或者用 reduction。7.3 什么时候该用融合算子生产环境里QKV 投影通常融合成一个矩阵乘法注意力分数和 softmax 也可能融合。llm.c 为了教学清晰可能分开写。你自己实现时先保证正确再考虑融合。融合能减少内存读写但调试难度大。7.4 关于数值精度的取舍GPT-2 原始是 float32但推理时可以用 float16 甚至 int8 量化。llm.c 如果做 float32精度没问题但内存占用大。如果你想省内存可以试试混合精度但要注意 LayerNorm 和 softmax 对精度敏感最好保持 float32。我在实际使用中发现LayerNorm 的eps取值对结果有影响1e-5是 GPT-2 的默认值但如果你自己训练可以试试1e-6或1e-4看哪个收敛更好。这个没有定论得实验。最后再分享一个小技巧如果你在调试注意力时发现输出全是 NaN先检查 softmax 有没有减最大值再检查掩码有没有用大负数最后检查缩放因子。这三个地方是 NaN 的高发区。按这个顺序排查基本能定位问题。