Transformer 这个架构从 2017 年那篇《Attention Is All You Need》发出来到现在已经快八年了。八年里我见过太多人读 Transformer读法基本分两派一派是抱着论文和图解硬啃把注意力公式背得滚瓜烂熟但一让他打开源码就懵另一派是直接 clone 一个实现跑起来能训能推理但问他为什么这里要除以根号 d_k、为什么用 LayerNorm 而不是 BatchNorm答不上来。这两派其实都没真正吃透它。我自己属于被逼着从源码里爬出来的那种。早些年做序列建模RNN 那套 LSTM、GRU 玩得挺熟后来要上一个长文本任务RNN 的串行瓶颈卡得死死的训练慢到怀疑人生。那时候硬着头皮去读 Transformer 源码一行一行对着论文抠才发现很多论文里一笔带过的东西源码里藏着真正的设计意图。这篇东西就是把我这些年读源码、改源码、踩坑的经验整理出来从张量形状怎么流动、每个模块为什么这么写、到实际改代码时哪些地方最容易翻车尽量讲透。适合已经看过论文但没读透源码的人也适合能跑通模型但想搞明白底层逻辑的人。读完你至少能做到拿到任何一份 Transformer 实现能快速定位每个模块在干什么能自己动手改结构、调参数而不是只会调包。1. 为什么必须结合源码读 Transformer1.1 论文和图解会骗你的三个地方先说个扎心的事实论文为了叙事流畅省略了大量工程细节而这些细节恰恰是理解模型行为的关键。我总结下来论文和图解最容易误导人的地方有三个。第一个是位置编码的实现方式。论文里给的公式是正弦余弦那一套很多人以为所有实现都这么干。但你打开源码会发现相当一部分实现用的是可学习的位置嵌入learnable positional embedding直接 nn.Embedding 一个 max_len 的查表。这两种方式在行为上差别不小正弦编码外推性好理论上能处理比训练时更长的序列可学习嵌入在训练长度内拟合得更充分但超出长度就抓瞎。论文只讲了前者源码里两种都有你不看源码根本不知道手上这份用的是哪种。第二个是残差连接和 LayerNorm 的顺序。论文里画的是 Post-LN也就是先做子层、再残差、最后 LayerNorm。但后来大家发现 Pre-LN先 LayerNorm 再进子层训练更稳尤其是深层模型。现在很多开源实现默认就是 Pre-LN跟论文图长得不一样。你要是照着论文图去理解源码看到 LayerNorm 在子层前面会以为自己看错了。第三个是注意力的 mask 机制。论文里对 mask 的描述非常简略就一句“我们用了 mask 防止看到未来”。但源码里 mask 的实现五花八门有的是加负无穷再 softmax有的是直接置零有的在 attention score 上加有的在 softmax 之后乘。不同实现方式在数值稳定性、梯度行为上都有差异。这些细节论文不会告诉你只有源码会。提示读 Transformer 源码时永远先确认三件事——位置编码是哪种、LN 放在哪、mask 怎么加的。这三件事决定了这份实现和论文的偏离程度。1.2 源码里藏着真正的设计权衡论文是结果导向的它告诉你“我们这么设计效果最好”但不告诉你“我们试过哪些别的、为什么放弃了”。源码不一样源码里的每一个选择背后都有工程上的权衡。举个例子多头注意力里头的 head 拆分。论文说把 d_model 拆成 h 个头每个头 d_k d_model / h。听起来很简单但源码实现上有两种主流写法一种是先做一次大的线性投影得到 q、k、v然后 reshape 成多头另一种是每个头独立做线性投影。前者计算效率高一次矩阵乘法搞定后者代码清晰但慢。绝大多数高效实现都用前者你去看源码里的 reshape 和 transpose 操作就能明白它是怎么把 (batch, seq, d_model) 变成 (batch, head, seq, d_k) 的。这个 reshape 的顺序如果搞错模型照样能跑但注意力就学废了而且不报错非常隐蔽。再比如前馈网络FFN的中间维度。论文里写的是 4 倍 d_model但这个 4 倍不是拍脑袋来的。源码里你会看到 d_ff 4 * d_model 这个设定它其实是个经验值后来有些实现改成 8/3 倍再对齐到 64 的倍数为了在特定硬件上更高效。这些变体论文不会提但源码里到处都是。1.3 读源码能帮你避开哪些实际坑我踩过的最大的坑是维度对不上但模型不报错。Transformer 里大量使用广播broadcasting和 reshape很多时候你的张量形状错了但因为广播机制代码照样跑loss 照样降只是效果差一截。你不读源码根本发现不了。还有一个坑是mask 的 dtype。mask 通常是 bool 或者 0/1 的 float但如果你在加负无穷的时候用了 float16-1e9 可能溢出成 -inf 再变成 nan训练直接崩。这种问题只有你盯着源码里 mask 那几行看才能提前预防。读源码的第三个好处是你能知道哪些参数是真的重要。比如 warmup 步数、学习率峰值、dropout 位置这些在论文里就是一句话但源码里你能看到它们具体作用在哪个张量上调的时候心里有数。2. 从张量形状入手拆解整体结构2.1 一次前向传播的张量流动全记录理解 Transformer 最快的方式不是看模块图而是跟踪一个 batch 的数据从头到尾形状怎么变。我拿一个具体例子走一遍假设 batch_size2序列长度 seq_len5d_model512头数 h8d_k64词表大小 vocab10000。输入是一串 token id形状 (2, 5)。第一步过 embedding变成 (2, 5, 512)。同时位置编码也是 (5, 512) 或者 (1, 5, 512)加进去形状不变。这时候你得到的是模型的输入表示。进入 encoder 层。先做多头注意力q、k、v 都是通过线性层从 (2, 5, 512) 投影得到形状还是 (2, 5, 512)。然后 reshape 成 (2, 5, 8, 64)再 transpose 成 (2, 8, 5, 64)。这时候 q 乘 k 的转置得到 attention score (2, 8, 5, 5)。除以根号 64 等于 8加 masksoftmax再乘 v得到 (2, 8, 5, 64)。然后 transpose 回来 reshape 成 (2, 5, 512)过输出线性层还是 (2, 5, 512)。残差加 LayerNorm形状不变。接着 FFN先升维到 (2, 5, 2048)激活再降回 (2, 5, 512)。残差加 LayerNorm。一层 encoder 走完形状始终是 (2, 5, 512)。堆 N 层形状不变。Decoder 稍微复杂点因为它有两个注意力子层。第一个是自注意力带 causal mask形状流动和 encoder 一样。第二个是交叉注意力q 来自 decoder (2, 5, 512)k、v 来自 encoder 输出 (2, 5, 512)attention score 是 (2, 8, 5, 5)。最后输出过线性层映射到 vocab变成 (2, 5, 10000)。把这条链路在纸上画一遍比看十张模块图都管用。你会发现整个 Transformer 的形状设计非常规整d_model 这条主线贯穿始终只有 FFN 中间和最后的输出层会变。2.2 为什么 d_model 要能被 head 整除d_model512h8d_k64整除得干干净净。这不是巧合是硬性要求。因为多头注意力的实现是把 d_model 维切成 h 份每份 d_k。如果 d_model 不能被 h 整除reshape 就会失败。但这里有个细节有些实现不要求整除它用 d_k d_model // h余数丢掉或者补零。这种做法很少见因为会浪费维度。主流实现都要求整除而且 d_k 最好是 2 的幂或者 64 的倍数这样 GPU 上的矩阵乘法效率最高。我实测过d_model512、h8 和 d_model512、h16d_k32在同样参数量下效果差别不大但 h16 的注意力头更“碎”每个头学到的模式更局部。h8 相对均衡。这个经验论文不会写但你改源码的时候可以自己试。2.3 位置编码到底加在哪一步位置编码加在 embedding 之后、进 encoder 之前。这个位置很关键因为一旦进了注意力层序列里每个位置的信息就开始互相混合了你必须在混合之前把位置信息注入进去。源码里通常是这么写的先算 embedding再算 positional encoding两者相加。注意是相加不是拼接。为什么相加因为拼接会让 d_model 翻倍参数量和计算量都上去了而且实验证明相加效果就够好。相加的隐含假设是embedding 空间和位置空间可以共享同一组维度模型自己会学会区分哪些维度表示语义、哪些表示位置。有个坑位置编码的 max_len 要设够。如果你训练时序列最长 512推理时来了个 600 的序列可学习位置嵌入就会越界报错。正弦编码没这个问题但很多实现为了简单还是用可学习嵌入。你读源码的时候一定要看 max_len 设了多少以及有没有做长度检查。3. 多头注意力源码里最值得抠的模块3.1 q、k、v 的投影为什么可以合并论文里 q、k、v 是三个独立的线性变换源码里你经常看到一个大矩阵一次性投影出三份。比如 d_model512它会定义一个 (512, 1536) 的权重一次矩阵乘法得到 (batch, seq, 1536)然后切成三份各 512。这么做纯粹是为了效率。三次小矩阵乘法不如一次大矩阵乘法快因为 GPU 对大矩阵乘法的利用率更高而且减少 kernel launch 次数。数学上完全等价因为三个独立的线性层拼在一起就是块对角矩阵等价于一个大矩阵。但这里有个实现细节切分的时候是按最后一维切还是按别的维度切。正确做法是 view 成 (batch, seq, 3, h, d_k)然后 permute 成 (3, batch, h, seq, d_k)再取 [0]、[1]、[2] 得到 q、k、v。这个 permute 的顺序如果写错q、k、v 就串了模型能跑但学不到东西。我见过有人把 3 放在最后切结果 q 里混了 k 的部分训练 loss 降得特别慢查了半天才发现是切分维度错了。3.2 缩放因子根号 d_k 的数学推导为什么是除以根号 d_k不是 d_k也不是别的这个推导其实不难但很多人没认真想过。假设 q 和 k 的每个分量都是独立同分布均值 0、方差 1。那么 q 和 k 的点积是 d_k 个乘积之和每个乘积的方差是 1所以点积的方差是 d_k。当 d_k64 时点积的标准差是 8。如果不缩放点积的值会落在正负几十的范围内。softmax 对大的输入非常敏感输入稍微大一点输出就接近 one-hot梯度几乎为零训练就停滞了。除以根号 d_k相当于把点积的方差拉回 1softmax 的输入落在合理范围内梯度健康。这就是缩放的全部意义。源码里通常写成q k.transpose(-2, -1) / math.sqrt(d_k)或者预先算好 scale 1 / math.sqrt(d_k) 再乘。注意有些实现在 d_k 很小时比如 16省略缩放因为影响不大。但只要 d_k 超过 32缩放就很有必要。你读源码时看到没有缩放的要警惕。3.3 mask 的三种实现方式与数值陷阱mask 是 Transformer 里最容易出数值问题的地方。我见过三种主流实现。第一种是加负无穷。构造一个 mask 矩阵被遮的位置是 0不遮的位置是 -inf然后加到 attention score 上。softmax 之后-inf 位置变成 0。这种写法最直观但有个陷阱如果 mask 和 score 的 dtype 不一致或者用了 float16-inf 可能变成 nan。稳妥的做法是用一个很大的负数比如 -1e9而不是真的 -inf。第二种是置零。先把 score 算出来被遮的位置直接赋 0然后再 softmax。这种写法的问题是softmax 之后被遮位置不是 0而是均匀分布的一小份因为 exp(0)1。所以必须在 softmax 之后再乘一次 mask把那些位置清零。两步操作容易漏。第三种是布尔 mask 配合 masked_fill。PyTorch 里用score.masked_fill(mask 0, -1e9)这是最干净的写法推荐。causal mask 的构造也有讲究。它是一个下三角矩阵位置 i 只能看到 j i。源码里通常用torch.tril(torch.ones(seq_len, seq_len))生成然后扩展成 (1, 1, seq_len, seq_len) 广播到所有 batch 和 head。3.4 多头拼接后为什么还要一个输出投影多头注意力算完之后把 h 个头的输出拼起来形状从 (batch, h, seq, d_k) 变回 (batch, seq, d_model)。这时候还要过一个线性层 W_o形状 (d_model, d_model)。为什么拼接之后还要投影因为拼接只是把不同头的信息放在不同维度上头与头之间没有交互。加一个输出投影相当于让不同头的信息在这个线性层里混合模型可以学习如何组合各个头的输出。如果没有这个投影各个头就是完全独立的表达能力受限。这个投影的参数量是 d_model 的平方是注意力模块里除了 qkv 投影之外最大的一块。你读源码时看到self.out_proj nn.Linear(d_model, d_model)就是它。4. 前馈网络与归一化的源码细节4.1 FFN 的 4 倍扩张从哪来FFN 的结构是 Linear(d_model, d_ff) - 激活 - Linear(d_ff, d_model)。d_ff 通常是 4 * d_model。这个 4 倍是论文里的设定但为什么是 4从直觉上理解FFN 是给模型提供非线性变换能力的地方。注意力层主要是加权求和本质是线性的softmax 那点非线性很弱。FFN 用两层线性加一个激活把每个位置独立地映射到高维再降回来相当于给每个位置做了一次特征变换。4 倍是个经验值太小了表达能力不够太大了参数爆炸。后来有研究说 8/3 倍配合 SwiGLU 激活效果更好但那是后话。源码里你会看到两种激活ReLU 和 GELU。原论文用 ReLU但后来的实现尤其是 BERT 之后大多用 GELU。GELU 更平滑梯度行为更好代价是计算稍慢。你读源码时看用的是哪个这会影响你复现的效果。4.2 LayerNorm 放在残差前还是后这是 Transformer 源码里最值得注意的差异之一。Post-LN 是LayerNorm(x Sublayer(x))Pre-LN 是x Sublayer(LayerNorm(x))。Post-LN 是原论文的写法但训练深层模型时不稳定需要 warmup 来救。Pre-LN 把 LayerNorm 提到子层前面梯度可以直接通过残差路径回传训练稳定得多现在几乎是默认选择。你打开一份现代实现大概率看到的是 Pre-LN。这个差异对复现影响很大。如果你用 Pre-LN 的代码去复现论文的 Post-LN 结果可能对不上。反过来也一样。读源码时先确认这一点能省很多调试时间。4.3 残差连接为什么能救深层网络残差连接x Sublayer(x)看起来简单但它是深层 Transformer 能训起来的关键。没有残差梯度在反向传播时要穿过每一层的非线性变换很容易消失或爆炸。有了残差梯度有一条高速公路直接回传即使子层学得不好至少 x 能原样传下去。源码里残差就是一行加法但要注意 dtype 和形状必须完全一致。有时候子层输出因为 dropout 或者别的操作形状变了加法就会广播广播错了不报错但结果不对。我习惯在残差加法前打印一下两个张量的 shape确认一致。5. 实操从零复现一个最小 Transformer5.1 环境与依赖准备我用的环境是 Python 3.10、PyTorch 2.1、CUDA 12.1。这些版本不是必须的但 PyTorch 2.x 的编译优化对 Transformer 提速明显建议用 2.0 以上。依赖就一个 torch不需要别的。pip install torch2.1.0如果你没有 GPUCPU 也能跑只是慢。最小 Transformer 参数量控制在几百万CPU 跑小 batch 没问题。5.2 核心模块的代码实现先写多头注意力。关键点是 qkv 合并投影、reshape 顺序、缩放、mask。import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.qkv_proj nn.Linear(d_model, 3 * d_model) self.out_proj nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) self.scale 1.0 / math.sqrt(self.d_k) def forward(self, x, maskNone): batch, seq, _ x.shape qkv self.qkv_proj(x) qkv qkv.view(batch, seq, 3, self.num_heads, self.d_k) qkv qkv.permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] score torch.matmul(q, k.transpose(-2, -1)) * self.scale if mask is not None: score score.masked_fill(mask 0, -1e9) attn torch.softmax(score, dim-1) attn self.dropout(attn) out torch.matmul(attn, v) out out.transpose(1, 2).contiguous().view(batch, seq, self.d_model) return self.out_proj(out)这段代码里view和permute的顺序是核心。view(batch, seq, 3, heads, d_k)把最后一维切成 3 份每份 heads*d_k。然后permute(2, 0, 3, 1, 4)把 3 提到最前面方便取 q、k、v。这个顺序错了qkv 就串了。接着写 FFN 和 encoder 层。class FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) self.act nn.GELU() def forward(self, x): return self.linear2(self.dropout(self.act(self.linear1(x)))) class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.attn MultiHeadAttention(d_model, num_heads, dropout) self.ffn FeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x, maskNone): x x self.dropout1(self.attn(self.norm1(x), mask)) x x self.dropout2(self.ffn(self.norm2(x))) return x注意这里用的是 Pre-LNLayerNorm 在子层前面。如果你想复现原论文把 norm 挪到加法后面。5.3 参数初始化与训练稳定性Transformer 的初始化很讲究。原论文用的是 Xavier 均匀初始化但后来大家发现对于深层模型把残差路径上的输出投影初始化得小一点比如乘以 1/sqrt(2N)N 是层数能显著提升稳定性。def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.LayerNorm): nn.init.ones_(m.weight) nn.init.zeros_(m.bias)我实测下来不做特殊初始化6 层以内的模型能正常训。超过 12 层不缩放残差输出loss 很容易在前期震荡。这个经验在论文里找不到是改源码改出来的。5.4 完整前向传播的验证方法写完模型别急着训。先做形状验证构造一个随机输入跑一遍前向打印每一层的输出形状确认和预期一致。然后做 mask 验证构造一个 causal mask检查 softmax 之后上三角是不是全零。最后做梯度验证跑一次反向检查每个参数的梯度是不是非零、有没有 nan。这三步做完基本能排除 90% 的低级错误。我见过太多人跳过验证直接训训了两天才发现 mask 写反了。6. 常见问题与排查技巧实录6.1 loss 不降或震荡的排查顺序遇到 loss 不降按这个顺序查先看学习率和 warmupTransformer 对 warmup 很敏感没有 warmup 前期容易崩再看 maskmask 写反了模型看到未来信息训练 loss 会异常低但验证集崩然后看 LayerNorm 的位置Post-LN 深层模型不 warmup 很难训最后看初始化残差输出没缩放的话深层模型前期震荡。我整理了一个速查表现象可能原因排查方法loss 一直不降学习率太小或 mask 全遮打印 mask 和 lrloss 前期震荡缺 warmup 或初始化太大加 warmup缩放残差输出loss 降但验证崩mask 写反看到未来检查 causal mask 方向loss 变 nanmask 用了 -inf 且 float16改用 -1e9训练极慢qkv 没合并投影合并成一次矩阵乘法6.2 维度不匹配的隐蔽 bug最隐蔽的 bug 是形状能广播但语义错了。比如 attention score 是 (batch, heads, seq, seq)mask 是 (batch, 1, seq, seq)广播没问题。但如果 mask 是 (batch, seq, seq)广播到 heads 维度时可能出错。我习惯在 mask 操作前打印 score 和 mask 的 shape确认广播维度对。还有一个是 reshape 和 view 的区别。view 要求内存连续permute 之后张量不连续直接 view 会报错要先 contiguous。这个错误很明显但如果你用 reshape 代替 view它内部会自动处理连续性反而掩盖了性能问题。我建议显式用 view contiguous心里有数。6.3 显存爆炸的几种典型场景Transformer 显存爆炸通常来自三个地方attention score 矩阵太大seq_len 平方FFN 中间维度太大以及没有用 gradient checkpointing。seq_len512 时score 是 512x512还好。seq_len4096 时score 是 4096x4096单个头就 64MB8 个头 512MBbatch 一大就爆。解决办法是 flash attention 或者分块计算但那是进阶话题。FFN 的 d_ff4d_model如果 d_model1024d_ff4096中间激活是 batchseq*4096也不小。可以适当降低 d_ff或者用 gradient checkpointing 换显存。提示显存不够时先降 batch size再降 seq_len最后才动模型维度。模型维度一改超参全要重调。6.4 推理阶段 kv cache 的源码实现推理时decoder 是自回归的每生成一个 token 就要重新算一遍注意力。如果不缓存 k、v每次都要把之前所有 token 重算复杂度是 O(n^2)。kv cache 就是把之前算过的 k、v 存下来新 token 只算自己的 q然后和缓存的 k、v 做注意力。源码里 kv cache 的实现是维护一个形状 (batch, heads, cache_len, d_k) 的张量每次新 token 来了把新的 k、v 拼接到 cache 后面。这样每步只算一个新位置的注意力复杂度降到 O(n)。这个机制在训练时不用只在推理时开。你读推理代码时看到 cache 相关的逻辑就是它。实现时要注意 cache 的长度管理和 mask 的配合新 token 的 mask 要能覆盖整个 cache。7. 从源码出发的进阶改造方向7.1 换激活函数与归一化层的实验读完源码最自然的改造就是换组件。把 GELU 换成 SwiGLU把 LayerNorm 换成 RMSNorm都是现在常见的做法。SwiGLU 需要三个线性层而不是两个参数量会变d_ff 要相应调整。RMSNorm 去掉了均值的计算只做方差归一化速度更快效果在多数任务上不输 LayerNorm。我做过对比同样参数量下SwiGLU RMSNorm 比 GELU LayerNorm 在语言建模任务上 perplexity 低一点点但训练速度快 10% 左右。这个收益在大模型上更明显。7.2 稀疏注意力与长序列适配标准注意力的 O(n^2) 复杂度是长序列的瓶颈。源码层面你可以改成滑动窗口注意力每个 token 只看前后 w 个或者分块注意力。改动点在 attention score 的 mask 上把全连接 mask 换成带状 mask。这个改造不难但要注意 mask 的构造和 cache 的配合。滑动窗口下kv cache 只需要存最近 w 个 token显存大幅下降。7.3 混合精度训练的源码改动点混合精度训练能省显存、提速但 Transformer 里有几个地方要小心。LayerNorm 通常要保持 float32softmax 也是。mask 的 -1e9 在 float16 下会溢出要用 -65504 或者保持 float32。源码改动主要是加 autocast 上下文然后对 LayerNorm 和 softmax 做特殊处理。PyTorch 的 amp 模块能自动处理大部分但 mask 那块要手动确认 dtype。我自己在实际操作中的体会是读 Transformer 源码这件事第一遍会很痛苦因为张量形状变来变去reshape 和 permute 看得头晕。但只要你耐着性子把一次前向传播的形状流动在纸上画一遍后面就顺了。真正值钱的不是记住每个模块怎么写而是理解每个设计选择背后的权衡——为什么这里要缩放、为什么那里要残差、为什么 mask 要这么加。这些理解到位了你改模型、调参数、排查 bug 的时候心里就有底而不是靠试。源码不会骗你它只是需要你愿意花时间去看。