★ FEATURED ARTICLE 1. 从 llm.c 小书第三篇说起为什么我们要死磕 GPT-2 的每一行很多人第一次看到 llm.c 这个项目时的反应都差不多一个用纯 C 语言写的、不依赖 PyTorch 的大模型训练与推理实现代码量不大却能把 GPT-2 完整跑起来。前两篇小书我们聊了环境搭建、张量内存布局、前向传播的整体骨架到了第三篇就绕不开整个 Transformer 里最“玄学”也最核心的部分——注意力机制以及它周围那几块看似不起眼、实则决定训练能否收敛的组件LayerNorm和MLP。我写这一系列笔记的初衷很简单现在大多数人调模型都是import torch然后model.train()一旦 loss 不降就抓瞎根本不知道梯度在内部经历了什么。而 llm.c 这种“把每一层都摊开给你看”的实现恰好是补课的最佳材料。GPT-2 的结构其实不复杂堆叠的 decoder block 里就三样东西多头自注意力、LayerNorm、前馈网络也就是 MLP。把这三块吃透你再看任何现代大模型的结构图都会有一种“哦还是这套东西换了个马甲”的淡定。这篇内容适合两类人一类是已经能跑通 llm.c、但对着attention那几十行 C 代码发懵的实践者另一类是想从零理解 GPT-2 内部机理、又不想被框架封装挡住视线的好奇者。我会把多头注意力机制的原理、LayerNorm 的归一化细节、MLP 的通道扩张逻辑以及它们在 C 代码里怎么落地一层层拆开讲。中间会穿插我自己踩过的坑比如为什么 LayerNorm 的 epsilon 不能乱设、为什么注意力缩放因子是1/sqrt(head_dim)而不是别的、为什么 MLP 要先升维再降维。这些细节才是决定你“跑得起来”还是“跑得对”的分水岭。2. 多头自注意力机制原理把一句话拆成多个视角去看2.1 从单头到多头为什么要“多此一举”先说单头注意力。给定一个序列的输入表示我们通过三个线性变换得到 Query、Key、Value 三组向量。核心动作就一个用 Query 和所有 Key 做点积算出每个位置对当前位置的“关注程度”softmax 归一化后再对 Value 加权求和。用一句话概括就是——当前位置的输出是所有位置信息按相关性强弱加权混合的结果。那为什么还要多头我一开始也觉得这是不是有点“为了复杂而复杂”。后来在 llm.c 里逐行对照才想明白单个注意力头只能学到一种“相关性模式”。比如一个头可能专注于捕捉主谓关系另一个头可能专注于捕捉指代关系还有一个头可能在看相邻词。如果只有一个头softmax 之后只有一组权重分布模型就被迫把所有关系压缩到同一个表示空间里表达能力受限。多头自注意力机制的做法是把embedding_dim切成num_heads份每一份独立做一次注意力计算最后把结果拼接起来再过一个输出投影。这样每个头在自己的子空间里学习不同的关注模式模型整体的表达能力就上来了。GPT-2 small 用的是 12 个头每个头维度 64总 embedding 维度 768。这个 768 12 × 64 的关系在 C 代码里体现为内存布局上的切分非常关键。2.2 缩放因子 1/sqrt(d_k)一个容易被忽略的数值稳定细节点积注意力的公式里有个缩放因子写出来是Attention(Q, K, V) softmax(Q·K^T / sqrt(d_k)) · V很多人背公式时会把sqrt(d_k)直接略过觉得它无关紧要。但我在 llm.c 里做数值对照实验时发现去掉这个缩放softmax 的输出会变得极其“尖锐”——几乎所有的概率质量都集中在一个位置上梯度接近于零训练直接卡死。原因其实不复杂假设 Q 和 K 的每个分量都是均值 0、方差 1 的独立随机变量那么它们点积之后的方差就是d_k。当d_k 64时点积结果的量级大概在 ±8 左右再大一点就会到 ±20。softmax 对输入的尺度非常敏感输入值一大指数运算就会把大的更大、小的趋近于零分布退化成 one-hot。除以sqrt(d_k)相当于把方差重新拉回到 1 附近让 softmax 工作在一个“温和”的区间里。提示这个缩放因子不是可学习的参数是固定的。在 C 实现里通常预先算好1.0f / sqrtf(head_dim)存成一个常量避免每次前向都做一次开方。2.3 因果掩码GPT 为什么不能“偷看未来”GPT 是自回归模型预测第 t 个 token 时只能看到前 t 个位置不能看到后面的。这个约束在注意力里靠因果掩码实现在 softmax 之前把上三角部分即 j i 的位置的注意力分数设成一个极大的负数softmax 之后这些位置的概率就变成 0。llm.c 里这块处理得很直接通常是在计算完Q·K^T之后用一个循环把非法位置填成-INFINITY或者一个足够大的负数比如-1e9。这里有个坑如果你用-INFINITY某些平台上exp(-inf)会返回 NaN导致整个 softmax 崩掉。我实测下来用-1e9这种“足够负但有限”的值更稳因为exp(-1e9)在 float32 下就是 0不会产生 NaN。另一个细节是掩码的时机。必须在 softmax 之前加不能 softmax 之后再乘一个 0/1 矩阵。因为 softmax 是全局归一化的你先算完再置零分母里已经包含了未来位置的贡献结果就错了。这个错误很隐蔽loss 会降但降不到应有的水平很多人会误以为是学习率问题。3. LayerNorm 在 GPT-2 里的位置与实现细节3.1 Pre-LN 还是 Post-LN一个影响训练稳定性的选择原始 Transformer 论文用的是 Post-LN也就是LayerNorm(x Sublayer(x))残差相加之后再归一化。但 GPT-2 用的是Pre-LN即x Sublayer(LayerNorm(x))先归一化再进子层残差连接保持“干净”。这个改动看起来只是顺序调换实际影响很大。Post-LN 在深层网络里梯度容易爆炸或消失需要精心设计 warmup 策略才能训起来。Pre-LN 把归一化放在子层入口梯度可以通过残差路径直接回传训练稳定性明显更好对 warmup 的依赖也小很多。llm.c 里实现的是 Pre-LN你在看代码时会发现 LayerNorm 出现在注意力和 MLP 之前而不是之后。3.2 归一化的计算过程均值、方差、缩放平移LayerNorm 对每个 token 的 embedding 向量独立做归一化不跨 batch、不跨序列位置。具体步骤是计算该向量所有元素的均值mean计算方差var归一化(x - mean) / sqrt(var eps)缩放平移gamma * normalized beta其中gamma和beta是可学习参数维度等于 embedding 维度。eps是为了防止方差为 0 时除零GPT-2 里通常取1e-5。在 C 实现里这一步需要两次遍历第一次求均值和方差第二次做归一化和缩放。也可以一次遍历用 Welford 算法在线更新但 llm.c 为了代码清晰用的是两次遍历。这里要注意数值精度如果你用 float32 累加一个 768 维向量的平方和误差还可以接受但如果维度上千最好用 double 累加或者 Kahan 求和否则方差算出来偏差会累积。注意eps的位置是加在方差上不是加在标准差上。写成sqrt(var eps)和sqrt(var) eps是两回事后者在 var 很小时会产生巨大误差。我见过有人写错这一行训练初期 loss 直接飙到 NaN。3.3 为什么 LayerNorm 对 GPT-2 如此关键有人会问BatchNorm 不也能归一化吗为什么 Transformer 系列清一色用 LayerNorm核心原因是序列长度可变。BatchNorm 在 batch 维度上统计不同样本的序列长度不一样padding 会污染统计量而且推理时如果 batch size 为 1BatchNorm 的统计量根本没法算。LayerNorm 对每个样本独立处理与 batch 大小、序列长度都无关天然适配 NLP 场景。另外LayerNorm 的缩放平移参数让网络在归一化之后仍能恢复一定的表达能力。如果只有归一化没有 gamma/beta相当于强制每层输出都是标准正态会损失信息。这两个参数虽然少但不可或缺。4. MLP 前馈网络注意力之后的“思考”环节4.1 先升维再降维4 倍扩张的由来GPT-2 每个 block 里注意力之后接一个 MLP。结构是两层线性变换加一个激活函数MLP(x) Linear2(GELU(Linear1(x)))其中Linear1把维度从 768 升到 30724 倍Linear2再降回 768。这个 4 倍扩张比例是原始 Transformer 论文定下来的后续大部分模型都沿用。为什么是 4 倍业界的普遍解释是注意力负责在序列维度上混合信息MLP 负责在每个位置的特征维度上做非线性变换。升维到 4 倍相当于给模型一个更宽的特征空间去做非线性组合表达能力强很多如果保持同维MLP 的容量就不够模型性能会明显下降。我在 llm.c 里做过一个小实验把扩张比例从 4 改成 2参数量少了训练 loss 下降明显变慢最终收敛水平也差一截。改成 8 呢参数量翻倍收益却不成比例。所以 4 倍是个性价比很高的经验值。4.2 GELU 激活比 ReLU 更平滑的选择GPT-2 用的激活函数是 GELU不是 ReLU。GELU 的公式是x * Φ(x)其中 Φ 是标准正态分布的累积分布函数。实际实现里常用近似GELU(x) ≈ 0.5 * x * (1 tanh(sqrt(2/π) * (x 0.044715 * x^3)))相比 ReLU 在负半轴直接截断为 0GELU 在负半轴有一个平滑的过渡梯度不会突然消失。这个平滑性对深层网络的训练稳定性有帮助。在 C 里实现 GELUtanh和pow都是标准库函数但要注意pow(x, 3)比x*x*x慢不少llm.c 里一般直接写成乘法。提示如果你在移植到其他平台时发现没有tanh可以用(exp(2x)-1)/(exp(2x)1)替代但要注意 x 很大时exp会溢出需要做截断保护。4.3 MLP 与注意力的分工一个负责“看”一个负责“想”我习惯把注意力理解成“信息收集”MLP 理解成“信息加工”。注意力让每个位置从其他位置拿信息MLP 则对拿到的信息做深度处理。两者交替堆叠模型就能既看到全局上下文又有足够的非线性容量去建模复杂关系。从参数量看MLP 其实占了每个 block 的大头。以 GPT-2 small 为例注意力的 QKV 投影加输出投影大概是4 * 768 * 768 ≈ 2.36M参数而 MLP 是768*3072*2 ≈ 4.72M正好是注意力的两倍。所以你说 MLP 是“配角”那是低估它了它才是参数量的主力。5. 在 llm.c 里落地从公式到 C 代码的关键步骤5.1 内存布局为什么 head 的切分方式很重要llm.c 里张量通常是一维数组加形状描述。对于注意力Q、K、V 的形状是[batch, seq_len, num_heads, head_dim]。这个布局决定了你在做多头计算时怎么用指针偏移去访问每个头的数据。常见做法是把num_heads * head_dim展平然后第 h 个头的 Q 就是从h * head_dim开始、长度head_dim的一段。这样每个头的计算是连续的缓存友好。如果你把布局搞成[batch, num_heads, seq_len, head_dim]虽然逻辑上更直观但跨头访问时步长变大性能会受影响。llm.c 选择了前者牺牲一点可读性换性能这是 C 实现的典型取舍。5.2 前向传播的完整流程一个 block 的前向大致是保存输入x作为残差x LayerNorm(x)计算 Q、K、V 投影对每个头做缩放点积注意力加因果掩码拼接多头输出过输出投影x 残差 注意力输出保存x作为新残差x LayerNorm(x)过 MLP升维、GELU、降维x 残差 MLP输出每一步都要保存中间结果因为反向传播需要。llm.c 里会预先分配好这些中间缓冲区避免运行时频繁 malloc。这个“预分配”思路在 C 实现里很常见也是它比动态图框架快的原因之一。5.3 反向传播里最容易出错的地方反向传播这块我踩过的坑主要集中在 LayerNorm 和 softmax 的梯度上。LayerNorm 的反向涉及均值、方差、归一化值三者的梯度耦合公式比较长。一个常见的错误是忘记gamma的梯度或者把beta的梯度算错。我的建议是先用数值梯度检验有限差分验证一遍确认无误再上大规模训练。Softmax 的反向有个经典陷阱当某个位置的 softmax 输出接近 0 或 1 时梯度会非常小容易下溢。llm.c 里通常会在 softmax 前减去最大值max subtraction来提升数值稳定性这个技巧在反向时也要对应处理。注意因果掩码位置在反向传播时梯度必须为 0。如果你在前向用-1e9掩码反向时这些位置的梯度自然接近 0但最好显式置零避免累积误差。6. 常见问题与排查技巧实录6.1 训练 loss 不降或震荡怎么定位遇到 loss 不降我一般按这个顺序排查排查项可能原因验证方法学习率太大导致震荡太小导致不降试 1e-4、3e-4、1e-3 三档LayerNorm eps太小导致除零太大导致归一化失效打印方差最小值确认 eps 合理注意力缩放忘记除以 sqrt(d_k)检查 softmax 输出是否过于尖锐因果掩码掩码时机错误或值不当检查掩码后 softmax 上三角是否为 0梯度爆炸深层网络梯度累积打印每层梯度范数这个表是我自己调试时总结的基本能覆盖八成问题。剩下两成通常是数据问题或者初始化问题。6.2 数值精度问题float32 够不够GPT-2 small 用 float32 训练是够的但要注意几个敏感点LayerNorm 的方差累加、softmax 的指数运算、以及损失函数的 log。如果发现 loss 偶尔出现 NaN优先检查这几处。混合精度训练时LayerNorm 和 softmax 通常要保持 float32否则精度损失会很明显。6.3 多头注意力的头数选择头数不是越多越好。头数增加每个头的维度就减小单头的表达能力下降。GPT-2 small 用 12 头、每头 64 维是个平衡点。如果你把头数加到 16每头只有 48 维性能反而可能下降。经验法则是保持每头维度在 64 左右头数根据总维度调整。6.4 实操心得先小后大先单层后多层我强烈建议在调试 llm.c 时先用 1 层、小维度比如 128、短序列比如 16跑通确认前向反向都对再逐步放大。直接上完整 GPT-2 调试一旦出错你根本不知道是哪一层的问题。这个“缩小规模验证”的思路在实现任何复杂模型时都适用。另外把中间张量的统计量均值、方差、最大值、最小值打印出来和 PyTorch 的参考实现逐层对比是定位问题最快的方法。llm.c 本身不提供这些打印你需要自己加但这几行代码能省你几天时间。7. 注意力机制的变体与延伸思考7.1 从自注意力到通道注意力、时序注意力多头自注意力机制的核心思想是“按相关性加权”这个思想可以迁移到很多场景。通道注意力机制比如 SE 模块就是在通道维度上做类似的事给每个通道算一个权重然后加权。时序注意力机制则是在时间维度上做加权常用于序列预测。理解了 QKV 这套框架再看这些变体你会发现它们只是换了“在哪个维度上算相关性”。7.2 效率优化KV Cache 与稀疏注意力GPT-2 推理时如果每次都重新算所有位置的 K、V计算量会随序列长度平方增长。KV Cache 的思路是缓存已算过的 K、V新 token 只算自己的 Q然后和缓存的 K、V 做注意力。llm.c 的推理部分就实现了这个优化能显著降低长序列的延迟。稀疏注意力则是另一个方向不是所有位置都互相看只保留一部分连接。这在超长序列场景下很有用但会损失一部分表达能力需要根据任务权衡。7.3 我个人的一点体会写 llm.c 小书这一系列最大的收获不是学会了用 C 写 Transformer而是被迫把每个公式都推导了一遍、每个数值都验证了一遍。以前用框架时觉得理所当然的东西比如 LayerNorm 的 eps、注意力的缩放、GELU 的近似自己实现时才发现每一个都有讲究。这种“从抽象到具体”的过程对理解模型本质帮助极大。如果你也在啃 llm.c我的建议是别急着跑通就完事把每个中间结果打印出来和参考实现对一遍。对上了你是真懂了对不上那正是你该花时间的地方。注意力机制、LayerNorm、MLP 这三块吃透GPT-2 对你来说就不再是黑盒了。
阅读完成 · 觉得有帮助？
核心技术解析与工程实践)
