说起来挺矛盾Transformer 是现在大语言模型背后的绝对核心但很多朋友第一次打开它的源码时往往盯了十几分钟也不知道该往哪看。这不是你基础差而是 Transformer 的代码存在一个天然门槛——它的逻辑是“分块”的和你之前见过的 CNN 或 RNN 那种“一层层往下堆”的直观结构不太一样。我最早硬啃这篇论文和开源实现时也经历过“每行都认识连起来就失忆”的阶段。这篇文章就是想把那段弯路省掉带你用一种更聪明的顺序去拆解 Transformer 代码先看骨架再补血肉最后跑起来。不管你是准备面试、想自己写个小型翻译模型还是单纯想搞清楚 LLM 内部到底长什么样这套拆法都能让你少走很多冤枉路。1. 开箱前先看结构Transformer 的模块边界与数据流1.1 为什么直接读源码会劝退新手我见过太多人一上来就打开某个开源仓库的model.py然后从class Transformer(nn.Module)开始往下顺。这个做法的问题在于你很容易被大量的assert、类型标注、mask分支和一堆叫不出名字的辅助函数带偏。Transformer 的核心其实非常单纯它就是一个把“序列”映射到“序列”的加工流水线。读源码的正确姿势是先抓住这条流水线的两头输入是什么输出是什么中间要过哪几道门。真实场景里你输入给模型的也不是什么高深东西就是一段已经切好的 token 序号比如[12, 34, 56, 78]。输出则是每个位置上、每个候选词的概率分布。中间的那些 Attention、LayerNorm、Feed-Forward本质都是对这批数字做加工和精炼。你不需要一开始就抠每个公式的推导你需要的是在脑子里画出一条线token 进向量化加位置层层提炼最后映射成概率。1.2 用一张“流水线图”代替死记硬背我在看代码之前建议你先在草稿纸上写下这样一行token ids - Embedding - Positional Encoding - [Attention - AddNorm - FFN - AddNorm] x N - Final Norm - Linear - Logits这行看起来像流水账但它其实就是 Transformer 的全部骨架。方括号里的部分叫作 Transformer Block通常堆 6 层、12 层甚至更多。每一层做的都是同一件事用 Attention 让 token 之间互相“交换信息”再用 FFN 对每个位置单独做一轮非线性变换。整个结构里没有循环、没有递归所有的序列关系都靠位置编码和注意力矩阵来承载。你在看任何开源实现时先定位这六个节点Embedding、PositionalEncoding、MultiHeadAttention、FeedForward、LayerNorm、FinalLinear。只要能把它们对应到代码里的具体类或函数恭喜你整个模型你已经理解了七成。剩下的时间不过是看形形色色的工程化变体Pre-Norm 还是 Post-Norm、相对位置编码怎么替换绝对位置、Flash Attention 怎么优化算力。这些后面再聊但骨架永远不会变。1.3 准备好实践环境和“最小配置”我习惯用 PyTorch 来做代码解析因为它的自动求导能帮我们省去大量手写反向传播的麻烦。你的环境不需要多豪华一个能跑 PyTorch 的笔记本就够用。建议先把接下来要用的配置参数写成一个简单的 dict 或 dataclass方便随时改import torch import torch.nn as nn import math config { vocab_size: 10000, # 词表大小 d_model: 512, # 向量维度所有嵌入和隐藏层都用这个数 n_heads: 8, # 多头注意力头数 d_ff: 2048, # 前馈网络中间层维度 max_len: 128, # 最大序列长度 n_layers: 6, # 编码器/解码器堆叠层数 dropout: 0.1, # 随机失活比例 }这些数字不是随便拍的。d_model512、n_heads8是论文里的经典配置d_model必须能整除n_heads不然多头拆分时会报错。我见过很多人在这一步栽跟头原因是没意识到 512 / 8 64每个头分到 64 维这个 64 才是每个注意力头的实际工作维度。你可以先跑一个非常迷你的版本比如d_model32, n_heads4, vocab_size500这样训练起来飞快调试也容易得多。2. 亲手写出核心代码embedding、位置编码与注意力机制2.1 Embedding 和位置编码给 token 一个“可学习的坐标”第一步是把离散的 token 变成密集向量。这一步在 PyTorch 里就是一个nn.Embedding(vocab_size, d_model)它维护一张可学习的查找表每个 token 会查出一行向量。这一步代码量很少但要注意一个容易忽略的细节位置编码也是加到 embedding 上的而且两者维度必须对齐。Transformer 本身不关心顺序所以我们必须显式地把位置信息塞进去。经典实现是用正弦和余弦函数生成一个固定矩阵之所以不直接用可学习的绝对位置向量原因有两个一是正弦编码可以推广到比训练时更长的序列二是它在相对位置上具备一定的归纳偏置。代码写出来大概是这样class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len512, dropout0.1): super().__init__() self.dropout nn.Dropout(pdropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len).unsqueeze(1) # [max_len, 1] div_term torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) # 偶数列用正弦 pe[:, 1::2] torch.cos(position * div_term) # 奇数列用余弦 pe pe.unsqueeze(0) # [1, max_len, d_model] self.register_buffer(pe, pe) def forward(self, x): x x self.pe[:, :x.size(1), :] return self.dropout(x)div_term那一行看起来有点吓人其实它就是公式里的1 / 10000^(2i/d_model)换了个写法。register_buffer把矩阵注册成模型的一部分这样它能跟着模型搬到 GPU但不会被当作参数更新。2.2 缩放点积注意力为什么除以根号 d_k 那么关键注意力机制可以理解成一个“寻址”过程每个 token 给出一个查询向量Q然后去所有 token 的键向量K上比对相似度得到权重后对值向量V做加权求和。注意力的核心代码通常只有几行但每行都藏着设计意图def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights torch.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output, attn_weights除以sqrt(d_k)这个操作新手最容易忽略但它极其重要。如果d_k很大比如 64那么Q和K的点积结果范围会很大跑到 softmax 的饱和区梯度会变得非常小甚至消失。除以sqrt(d_k)是为了把方差拉回 1 附近让 softmax 的输入保持在一个合理范围内梯度才能顺利传播。类比一下就是你在黑漆漆的房间里找钥匙手里电筒的亮度必须调到合适档位太暗看不见太亮又会白茫茫一片什么都分不清。sqrt(d_k)就是这个调光旋钮。masked_fill(mask 0, -inf)也值得多看一眼。我经常看到有人误以为 mask 应该把目标位置填成 0然后以为 softmax 后权重自然为 0这是大错特错的。softmax 里填 0 会让位置分到一份本来不属于它的权重只有填入负无穷softmax 之后才会得到真正的 0。这一步在后面做解码器因果掩码时非常容易踩坑。2.3 多头注意力把认知拆分到多个“视角”多头注意力不是把注意力做 8 遍而是把d_model512的向量切成 8 份每份 64 维各自做缩放点积注意力最后拼回来。为什么要切你可以把每个头理解成不同的“观察角度”有的头更关注语法关系有的头更关注距离近的相邻词有的头则对长距离依赖更敏感。这也是 Transformer 表达能力的关键来源。代码上有一个顺手但必须注意的细节PyTorch 的nn.Linear默认喂入和输出的都是最后一维所以我们要先把多头拆分开做完注意力再合并回去。一个常见写法是class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0 self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads self.w_q nn.Linear(d_model, d_model) self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, query, key, value, maskNone): batch_size query.size(0) # 投影 Q self.w_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K self.w_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V self.w_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) attn_output, attn_weights scaled_dot_product_attention(Q, K, V, mask) attn_output attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.out_proj(attn_output), attn_weights那一串.view(...).transpose(...)就是多头拆分。你完全可以把它拆开打印每一步的 shape我强烈建议你这么做。我最早看这段代码时一直在脑子里玩魔方一样想维度变换后来干脆在 notebook 里打印[batch, seq, d_model] - [batch, heads, seq, d_k]一下就豁然开朗了。2.4 前馈网络与残差连接每个 token 的“独立思考”注意力层负责让 token 之间交流但只有交流还不够。前馈网络Feed-Forward Network的作用是让每个 token 在“听完别人的观点”之后独立地进行一轮深层加工。它其实就是两个线性层夹一个激活函数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) def forward(self, x): return self.linear2(self.dropout(torch.relu(self.linear1(x))))d_ff通常设为4 * d_model这里就是 2048。也就是说FFN 把 512 维先暴力扩张到 2048 维再压回 512 维。这种“先变胖再变瘦”的结构本质上是在给每个位置更多的参数容量去拟合复杂映射。加上残差连接后整个 Block 就是x x Dropout(FFN(LN(x)))这种做法在实现上叫 Pre-Norm即先归一化再做子层。相比老的 Post-NormPre-Norm 在深层网络中更稳定训练初期 loss 下降也更平滑现在的主流开源模型基本都采用它。3. 前向传播全记录从 token 到 logits 发生了什么3.1 组装一个最小可用的 Transformer Encoder我习惯先用 Encoder 来验证整套流程因为它的代码最短而且能直观地看到输出维度始终不变。下面是组装好的最小前向代码包含我们前面所有零件class TransformerEncoder(nn.Module): def __init__(self, config): super().__init__() self.embedding nn.Embedding(config[vocab_size], config[d_model]) self.pos_encoding PositionalEncoding(config[d_model], config[max_len], config[dropout]) self.layers nn.ModuleList([ TransformerBlock(config) for _ in range(config[n_layers]) ]) self.norm nn.LayerNorm(config[d_model]) def forward(self, src, src_maskNone): x self.embedding(src) * math.sqrt(config[d_model]) x self.pos_encoding(x) for layer in self.layers: x layer(x, src_mask) return self.norm(x)这里有一个细节self.embedding(src) * math.sqrt(d_model)。这一步在论文里写得很明白但很多代码注释都会略过。因为 embedding 是随机初始化的它的方差比较小乘上sqrt(d_model)是为了让 embedding 的数值范围和位置编码匹配防止位置编码把信号淹掉。这属于“不加也行但加了更稳”的工程细节。src_mask通常是 padding mask因为一个 batch 里的句子长度不一样短句子后面要用padtoken 补齐而这些 padding 位置不应该参与注意力。实现上mask 的形状一般是[batch, 1, 1, src_len]它会沿着注意力分数的形状做广播把 padding 位置全部屏蔽掉。3.2 TransformerBlock 的完整实现与形状追踪为了让打印形状变得更方便我在TransformerBlock里会故意保留几个临时变量方便在 notebook 里观察数据变换class TransformerBlock(nn.Module): def __init__(self, config): super().__init__() self.attn MultiHeadAttention(config[d_model], config[n_heads], config[dropout]) self.ffn FeedForward(config[d_model], config[d_ff], config[dropout]) self.norm1 nn.LayerNorm(config[d_model]) self.norm2 nn.LayerNorm(config[d_model]) self.dropout nn.Dropout(config[dropout]) def forward(self, x, maskNone): # 第一段多头注意力 残差 attn_out, _ self.attn(x, x, x, mask) x x self.dropout(attn_out) x self.norm1(x) # Pre-Norm先归一化 # 第二段前馈网络 残差 ffn_out self.ffn(self.norm2(x)) x x self.dropout(ffn_out) return x看起来顺序有些别扭常见的架构图是先画 AddNorm再进 FFN这里我在 FFN 之前又加了一次self.norm2(x)。这其实就是 Pre-Norm 的标准写法每个子层输入前先归一化残差连接在外面。你可以把它改造成 Post-Norm也就是归一化放到残差后面。两种都能跑但你会发现 Pre-Norm 对学习率更宽容在层数更深时表现更稳定。跑一次前向你会看到维度全程不变输入是[batch, seq_len, d_model]中间每一层的输出仍然是[batch, seq_len, d_model]。这就给了你一个巨大的调试便利无论模型中间哪一段出问题你只需要比较它的输入输出维度是否保持一致。我实际带过的一些朋友最初写代码时反复在 shape 上报错后来养成了“每层打印 shape”的习惯五分钟内就能定位问题。3.3 Decoder 的差异因果掩码与 Cross-AttentionEncoder 和 Decoder 的代码结构高度相似区别主要在两个地方。第一个是 Decoder 做的是“自回归”生成预测第 i 个 token 时只能看前 i-1 个 token不能偷看未来。这个限制靠因果掩码causal mask实现代码通常是causal_mask torch.tril(torch.ones(max_len, max_len)).view(1, 1, max_len, max_len).bool()torch.tril生成一个下三角矩阵矩阵里位置[i, j]为 1 表示第 i 个 token 能看到第 j 个 token当 j i 时。在注意力分数矩阵里这个 mask 会把它右上角的未来信息全部变成-inf。调试时最容易出的问题是 mask 维度没对齐scores的形状是[batch, heads, tgt_len, src_len]所以 mask 至少要能广播成这个形状建议维度是[batch, 1, tgt_len, src_len]或[1, 1, tgt_len, src_len]。第二个区别是 Decoder 里的注意力有两种第一层是自回归自注意力看的是自己之前生成的 token第二层是 Cross-AttentionQ来自 DecoderK和V来自 Encoder 的输出。这就是“翻译”的关键Decoder 从目标语言里拿到已经生成的部分同时从源语言里拿到已经编码好的完整语义。代码里Cross-Attention 其实不需要额外写一个新类直接复用MultiHeadAttention只要 forward 时传入不同的 K 和 V 就行enc_out encoder(src) dec_out decoder(tgt, enc_out, tgt_mask, src_mask)我个人做调试时会刻意把 Encoder 和 Decoder 的输出分开打印。因为在实际训练初期Decoder 的 loss 下降速度通常比 Encoder 慢得多如果一开始就合在一起训练你会很难判断是编码端没学好还是解码端没学好。虽然很多人直接端到端训练但对刚接触代码的人来说单独跑一遍两个模块能让你理解得更透。4. 训练与推理实战一个小模型的完整生命周期4.1 数据准备与 padding mask 的生成训练一个 mini Transformer我建议用一个小型平行语料比如英法简表、或者随便一对句子序列重点不是效果多好而是让你完整跑通流程。原始文本要先分词、转成 id、按 batch 拼装成定长张量。这里有一个经典问题每个 batch 是最长句子长度还是统一固定长度实践中我一般用动态 padding也就是每一个 batch 只 pad 到该 batch 内最长句子的长度。这样不会浪费大量算力在padtoken 上。生成 padding mask 的代码也很简单def make_pad_mask(seq, pad_idx0): # seq: [batch, seq_len] return (seq ! pad_idx).unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len]注意这个 mask 最终要变形到注意力分数能广播的形状。如果你的scores是[batch, heads, tgt_len, src_len]那[batch, 1, 1, src_len]就可以沿着heads和tgt_len广播。4.2 组合训练循环与损失计算训练循环写出来后你会发现 Transformer 和别的神经网络没有本质区别前向、算 loss、反向、更新。但有两个细节要专门说。第一交叉熵损失计算时nn.CrossEntropyLoss需要直接把 logits 扔进去它内部会做 softmax如果你在外面又做一次 softmax反而会算错梯度。第二padding 位置不能被算进 loss所以要把损失函数的ignore_index设为pad的 id。criterion nn.CrossEntropyLoss(ignore_index0) optimizer torch.optim.Adam(model.parameters(), lr1e-3, betas(0.9, 0.98), eps1e-9)betas(0.9, 0.98)是论文里推荐的超参它和我们平时常用的(0.9, 0.999)略有不同。更稳妥的做法是给 Adam 配一个学习率预热warmup前 4000 步让学习率从很小的值线性涨到峰值之后再按步数衰减。因为 Transformer 的初始化方差相对较大如果一开始就用大学习率容易出现早期 loss 震荡甚至直接发散。训练几个 epoch 后如果 loss 在稳定下降你会看到它从最初的 10 左右一路降到 2 以内。遇到 loss 突然变成 NaN 时先别急着调模型把学习率调低一个数量级再试。我见过大量 NaN 案例九成是学习率太高或注意力没有缩放只有一成是数据里有无穷值。4.3 推理阶段与贪心解码训练完毕推理阶段就是自回归式地一个词一个词往外蹦。核心逻辑是给定起始 tokenbos喂进 Decoder 得到下一个 token 的概率取概率最高的那个拼到序列后面再重复这个过程直到生成eos或者达到最大长度。伪代码如下def greedy_decode(model, encoder_output, src_mask, start_token, end_token, max_len50): tgt torch.tensor([[start_token]]).to(src_mask.device) for _ in range(max_len): tgt_mask make_causal_mask(tgt.size(1)).to(src_mask.device) logits model.decode(encoder_output, tgt, tgt_mask, src_mask) # [1, tgt_len, vocab] next_token logits[:, -1, :].argmax(dim-1) if next_token.item() end_token: break tgt torch.cat([tgt, next_token.unsqueeze(0)], dim-1) return tgt有一个很常见的 bug生成时忘记更新tgt_mask的长度。tgt每增加一个 token因果 mask 的尺寸也要跟着变否则新位置看不到旧位置或者根本找不到它的 mask直接报维度错误。这个问题我在自己写的第一个翻译 demo 里也踩过后来索性把 mask 生成函数做得更“自动化”一点直接根据tgt.size(1)动态生成省得每次手滑。4.4 训练日志怎么读loss、困惑度与过拟合信号看训练日志不能只看 loss还要结合困惑度perplexity来判断模型是不是真的在学。困惑度等于exp(loss)。如果 loss 到 2.0困惑度就是 7.4意思是模型平均在每个位置上对 7.4 个候选词感到困惑这算一个中等水平如果 loss 到 1.0困惑度约 2.7说明模型已经比较有把握了。对一个小型 demo 来说不要指望它达到大模型那种 1.0 以下的 loss更重要的指标是验证集 loss 是否和训练集 loss 同步下降。一旦发现训练 loss 还在降、验证 loss 开始回升恭喜你过拟合了。这时候应该加大 dropout、减小模型容量或增加数据量。5. 常见坑位排查注意力 NaN、mask 错位与头数不匹配5.1 shape 不匹配与维度混淆速查我把常见的 shape 错误整理成了一张速查表希望你下次遇到报错能直接对号入座报错信息常见原因解决办法Expected 4D input注意力打分时 Q 和 K 没有扩到 4 维确认view后已transpose(1,2)The size of tensor a (8) must match size of tensor b (16)n_heads与d_model不整除将d_model除以n_heads时确保是整数mask shape mismatchmask 维度不能广播到scores统一 mask 为[batch, 1, tgt_len, src_len]CUDA error: device-side assert triggered通常是ignore_index没设好或 token id 越界检查 label 里是否混入非法 idone element cannot be broadcastedEmbedding 与位置编码 seq 维度不匹配切取pe[:, :x.size(1)]头数不匹配这个坑我在很多小白代码里反复看到。d_model512n_heads6这时候 512 不能被 6 整除得到的d_k是 85.33你就没法用view拆分了。要么把n_heads改成 8要么把d_model改成 480 或 504。工程上很少为了迁就头数去改d_model基本都是把头数调到能整除的数值。5.2 注意力 NaN 的定位方法NaN 是最让人头大的问题因为它不会告诉你具体是哪一行。我的排查顺序是先看 loss 在哪一步变成 NaN再在 forward 里插入中间打印检查 attention scores 和 softmax 输出的数值范围。如果是scores爆炸检查有没有除以sqrt(d_k)如果是scores里出现-inf检查 mask 是不是把该看的 token 也屏蔽了。还有一个容易被忽略的源头FFN 里的 ReLU 输出范围太大。如果d_ff很大ReLU 之后有些位置的数值可能是几十甚至上百再加上残差连接数值一路滚雪球。这也是很多现代实现把 ReLU 换成 GELU 或 SwiGLU 的原因之一。GELU 在负半轴更平滑不会像 ReLU 那样直接一刀切。如果训练中早期 loss 一直上蹿下跳可以考虑把激活函数换成nn.GELU()。5.3 mask 错位的三类典型表现mask 错位通常有三种表现生成时出现重复词、训练 loss 比预期高、或者模型在测试时“偷看未来”导致生成的连贯性很差。我遇到最隐蔽的一种是mask 形状对了但数值反了。比如本该屏蔽的位置填了 1该保留的位置反而被屏蔽成了 0。判断技巧很简单把 mask 打印出来确认下三角矩阵的 1 在左下角而右上方是 0。拿tgt_len4来说合法 mask 应该是[[1, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]这个矩阵的生成方式有很多种但核心只有一条i行j列表示第i个输出 token 是否能看到第j个输入 token。只要这个逻辑搞对了mask 基本不会错。5.4 残差连接与 LayerNorm 的位置混淆Pre-Norm 和 Post-Norm 的实现差异虽然只有几行代码但对训练行为影响很大。Post-Norm 的写法更贴近原始论文x Norm(x Sublayer(x))。Pre-Norm 的写法则是x x Sublayer(Norm(x))。如果同时使用大学习率和 Post-Norm浅层容易出现梯度不稳定Pre-Norm 则相对友好。但 Pre-Norm 有一个副作用最后一层输出和残差直接相加导致输出数值可能没有被充分归一化。所以很多实现会在最后加一个Final LayerNorm。如果你搭建的模型最后一层后直接接nn.Linear到词表我的建议是加一层nn.LayerNorm(d_model)作为收尾能显著提升稳定性。这也是我在看很多开源代码时发现的一个共性设计不是论文里写的但确实是大家踩过坑后总结出的经验。6. 从代码到理解注意力可视化、参数量与扩展方向6.1 用一张注意力热力图验证模型真的在“学”模型跑起来以后很多人就想看一眼 attention 到底学到了什么。最简单的方法是让MultiHeadAttention的 forward 返回权重矩阵然后画热力图。你会看到不同头有不同的关注模式有些头对角分布说明它主要盯紧当前位置和前一个位置有些头有较宽的横向条纹说明它在做长距离的实体对齐。一个小技巧是把注意力权重和对应的句词打印出来对照着看。我曾经在一个中文翻译 demo 里发现某个头特别关注“你”和“我”这类代词另一个头则总是把注意力均匀撒在标点和停用词上。这种可视化并不仅仅是为了“好看”它能帮你确认训练是否正常。如果一个头在所有层里的注意力权重都是均匀分布可能性有两种要么还没收敛要么这个头已经“死”掉了。6.2 参数量估算为什么 512 维这么流行用手算一遍参数量能让你对模型规模有更具体的感知。Embedding 的参数量是vocab_size * d_model1 万词表加 512 维就是 512 万。多头注意力的三个投影加一个输出投影每个都是512*512总共约 104 万。FFN 的两个线性层则是512*2048 2048*512约 209 万。单层 Block 大约 300 多万参数6 层叠加后主体就是 2000 多万再加上 Embedding 与其他一个小型 Transformer 大概在 2500 万参数左右。这个规模在普通显卡上完全能训起来但如果你想带着 3 万词表跑光 Embedding 就快 1.5 亿参数内存压力就上来了。现在你也能理解为什么d_model512是经典之选它既能表达足够丰富的信息又让多头注意力和 FFN 的参数量维持在可控范围内。真的跑到 GPT 级别的模型d_model往往在 768 到 12288 之间参数量随维度平方级膨胀这也是为什么大模型动辄几百亿参数。6.3 为了让代码更现代Flash Attention、相对位置编码和 KV Cache如果你已经跑通基础版本我建议你接着研究三个扩展点。第一个是 Flash Attention它的核心不是改变注意力公式而是通过分块计算和重计算减少显存占用、提高速度。原来的注意力需要先把[seq_len, seq_len]的完整分数矩阵算出来并存到显存里Flash Attention 则把这块矩阵切成小块边算边用内存占用从平方级降为线性级。第二个是相对位置编码经典的正弦编码已经不太够用。很多新模型用旋转位置编码RoPE或 ALiBi它们不仅编码绝对位置还能让模型更好感知 token 之间的相对距离。第三个是 KV Cache在推理阶段把已经算过的K和V缓存起来避免每生成一个新 token 就把过去所有位置的K、V重算一遍。推理性能的提升非常明显尤其是长文本生成时几乎是从“不能忍”变成“丝滑流畅”。6.4 我的个人实践心得讲真Transformer 代码本身并不难难的是你容易被各种花哨的工程优化吓到。我见过太多人花了一个星期去追最新的稀疏注意力实现结果最后连最基础的scaled_dot_product_attention都写不顺。我的体会是先把最朴素的版本跑通再一步步往上加工程技巧这样你才能分清哪些修改是真正必要的哪些只是锦上添花。另外我很推荐你在调试时顺手加一个“形状断言工具”——可以是你自己的函数也可以是 PyTorch 内置的assert。每次输入输出经过某个子层时都检查维度是否符合预期。这个过程很笨但极其有效。我在实际项目里用它节省的时间远比写它用的时间多得多。一个小技巧是给每个子层起一个简短的名字比如emb、pos、attn、ffn然后打印形状时带着名字输出这样出错时一眼就能看见是哪段炸了。最后再分享一个小经验如果你真想彻底搞懂 Transformer不要只盯着别人的代码看挑一个周末自己从头写一遍写完再跑一遍机器翻译的 demo。这个过程的收获比你看十篇解析文章都大。写不出来也没关系卡住的地方恰恰是你理解最薄弱的地方把它们标记下来再回头看我前面讲的这些顺序去补课你会发现自己突然就能看懂那些开源巨作的源码了。 
长河君 · 2026/10/13 23:11:52 · 阅读 1.2k Transformer源码这样拆：先骨架后血肉，从零跑通核心代码
说起来挺矛盾Transformer 是现在大语言模型背后的绝对核心但很多朋友第一次打开它的源码时往往盯了十几分钟也不知道该往哪看。这不是你基础差而是 Transformer 的代码存在一个天然门槛——它的逻辑是“分块”的和你之前见过的 CNN 或 RNN 那种“一层层往下堆”的直观结构不太一样。我最早硬啃这篇论文和开源实现时也经历过“每行都认识连起来就失忆”的阶段。这篇文章就是想把那段弯路省掉带你用一种更聪明的顺序去拆解 Transformer 代码先看骨架再补血肉最后跑起来。不管你是准备面试、想自己写个小型翻译模型还是单纯想搞清楚 LLM 内部到底长什么样这套拆法都能让你少走很多冤枉路。1. 开箱前先看结构Transformer 的模块边界与数据流1.1 为什么直接读源码会劝退新手我见过太多人一上来就打开某个开源仓库的model.py然后从class Transformer(nn.Module)开始往下顺。这个做法的问题在于你很容易被大量的assert、类型标注、mask分支和一堆叫不出名字的辅助函数带偏。Transformer 的核心其实非常单纯它就是一个把“序列”映射到“序列”的加工流水线。读源码的正确姿势是先抓住这条流水线的两头输入是什么输出是什么中间要过哪几道门。真实场景里你输入给模型的也不是什么高深东西就是一段已经切好的 token 序号比如[12, 34, 56, 78]。输出则是每个位置上、每个候选词的概率分布。中间的那些 Attention、LayerNorm、Feed-Forward本质都是对这批数字做加工和精炼。你不需要一开始就抠每个公式的推导你需要的是在脑子里画出一条线token 进向量化加位置层层提炼最后映射成概率。1.2 用一张“流水线图”代替死记硬背我在看代码之前建议你先在草稿纸上写下这样一行token ids - Embedding - Positional Encoding - [Attention - AddNorm - FFN - AddNorm] x N - Final Norm - Linear - Logits这行看起来像流水账但它其实就是 Transformer 的全部骨架。方括号里的部分叫作 Transformer Block通常堆 6 层、12 层甚至更多。每一层做的都是同一件事用 Attention 让 token 之间互相“交换信息”再用 FFN 对每个位置单独做一轮非线性变换。整个结构里没有循环、没有递归所有的序列关系都靠位置编码和注意力矩阵来承载。你在看任何开源实现时先定位这六个节点Embedding、PositionalEncoding、MultiHeadAttention、FeedForward、LayerNorm、FinalLinear。只要能把它们对应到代码里的具体类或函数恭喜你整个模型你已经理解了七成。剩下的时间不过是看形形色色的工程化变体Pre-Norm 还是 Post-Norm、相对位置编码怎么替换绝对位置、Flash Attention 怎么优化算力。这些后面再聊但骨架永远不会变。1.3 准备好实践环境和“最小配置”我习惯用 PyTorch 来做代码解析因为它的自动求导能帮我们省去大量手写反向传播的麻烦。你的环境不需要多豪华一个能跑 PyTorch 的笔记本就够用。建议先把接下来要用的配置参数写成一个简单的 dict 或 dataclass方便随时改import torch import torch.nn as nn import math config { vocab_size: 10000, # 词表大小 d_model: 512, # 向量维度所有嵌入和隐藏层都用这个数 n_heads: 8, # 多头注意力头数 d_ff: 2048, # 前馈网络中间层维度 max_len: 128, # 最大序列长度 n_layers: 6, # 编码器/解码器堆叠层数 dropout: 0.1, # 随机失活比例 }这些数字不是随便拍的。d_model512、n_heads8是论文里的经典配置d_model必须能整除n_heads不然多头拆分时会报错。我见过很多人在这一步栽跟头原因是没意识到 512 / 8 64每个头分到 64 维这个 64 才是每个注意力头的实际工作维度。你可以先跑一个非常迷你的版本比如d_model32, n_heads4, vocab_size500这样训练起来飞快调试也容易得多。2. 亲手写出核心代码embedding、位置编码与注意力机制2.1 Embedding 和位置编码给 token 一个“可学习的坐标”第一步是把离散的 token 变成密集向量。这一步在 PyTorch 里就是一个nn.Embedding(vocab_size, d_model)它维护一张可学习的查找表每个 token 会查出一行向量。这一步代码量很少但要注意一个容易忽略的细节位置编码也是加到 embedding 上的而且两者维度必须对齐。Transformer 本身不关心顺序所以我们必须显式地把位置信息塞进去。经典实现是用正弦和余弦函数生成一个固定矩阵之所以不直接用可学习的绝对位置向量原因有两个一是正弦编码可以推广到比训练时更长的序列二是它在相对位置上具备一定的归纳偏置。代码写出来大概是这样class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len512, dropout0.1): super().__init__() self.dropout nn.Dropout(pdropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len).unsqueeze(1) # [max_len, 1] div_term torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) # 偶数列用正弦 pe[:, 1::2] torch.cos(position * div_term) # 奇数列用余弦 pe pe.unsqueeze(0) # [1, max_len, d_model] self.register_buffer(pe, pe) def forward(self, x): x x self.pe[:, :x.size(1), :] return self.dropout(x)div_term那一行看起来有点吓人其实它就是公式里的1 / 10000^(2i/d_model)换了个写法。register_buffer把矩阵注册成模型的一部分这样它能跟着模型搬到 GPU但不会被当作参数更新。2.2 缩放点积注意力为什么除以根号 d_k 那么关键注意力机制可以理解成一个“寻址”过程每个 token 给出一个查询向量Q然后去所有 token 的键向量K上比对相似度得到权重后对值向量V做加权求和。注意力的核心代码通常只有几行但每行都藏着设计意图def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights torch.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output, attn_weights除以sqrt(d_k)这个操作新手最容易忽略但它极其重要。如果d_k很大比如 64那么Q和K的点积结果范围会很大跑到 softmax 的饱和区梯度会变得非常小甚至消失。除以sqrt(d_k)是为了把方差拉回 1 附近让 softmax 的输入保持在一个合理范围内梯度才能顺利传播。类比一下就是你在黑漆漆的房间里找钥匙手里电筒的亮度必须调到合适档位太暗看不见太亮又会白茫茫一片什么都分不清。sqrt(d_k)就是这个调光旋钮。masked_fill(mask 0, -inf)也值得多看一眼。我经常看到有人误以为 mask 应该把目标位置填成 0然后以为 softmax 后权重自然为 0这是大错特错的。softmax 里填 0 会让位置分到一份本来不属于它的权重只有填入负无穷softmax 之后才会得到真正的 0。这一步在后面做解码器因果掩码时非常容易踩坑。2.3 多头注意力把认知拆分到多个“视角”多头注意力不是把注意力做 8 遍而是把d_model512的向量切成 8 份每份 64 维各自做缩放点积注意力最后拼回来。为什么要切你可以把每个头理解成不同的“观察角度”有的头更关注语法关系有的头更关注距离近的相邻词有的头则对长距离依赖更敏感。这也是 Transformer 表达能力的关键来源。代码上有一个顺手但必须注意的细节PyTorch 的nn.Linear默认喂入和输出的都是最后一维所以我们要先把多头拆分开做完注意力再合并回去。一个常见写法是class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0 self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads self.w_q nn.Linear(d_model, d_model) self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, query, key, value, maskNone): batch_size query.size(0) # 投影 Q self.w_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K self.w_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V self.w_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) attn_output, attn_weights scaled_dot_product_attention(Q, K, V, mask) attn_output attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.out_proj(attn_output), attn_weights那一串.view(...).transpose(...)就是多头拆分。你完全可以把它拆开打印每一步的 shape我强烈建议你这么做。我最早看这段代码时一直在脑子里玩魔方一样想维度变换后来干脆在 notebook 里打印[batch, seq, d_model] - [batch, heads, seq, d_k]一下就豁然开朗了。2.4 前馈网络与残差连接每个 token 的“独立思考”注意力层负责让 token 之间交流但只有交流还不够。前馈网络Feed-Forward Network的作用是让每个 token 在“听完别人的观点”之后独立地进行一轮深层加工。它其实就是两个线性层夹一个激活函数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) def forward(self, x): return self.linear2(self.dropout(torch.relu(self.linear1(x))))d_ff通常设为4 * d_model这里就是 2048。也就是说FFN 把 512 维先暴力扩张到 2048 维再压回 512 维。这种“先变胖再变瘦”的结构本质上是在给每个位置更多的参数容量去拟合复杂映射。加上残差连接后整个 Block 就是x x Dropout(FFN(LN(x)))这种做法在实现上叫 Pre-Norm即先归一化再做子层。相比老的 Post-NormPre-Norm 在深层网络中更稳定训练初期 loss 下降也更平滑现在的主流开源模型基本都采用它。3. 前向传播全记录从 token 到 logits 发生了什么3.1 组装一个最小可用的 Transformer Encoder我习惯先用 Encoder 来验证整套流程因为它的代码最短而且能直观地看到输出维度始终不变。下面是组装好的最小前向代码包含我们前面所有零件class TransformerEncoder(nn.Module): def __init__(self, config): super().__init__() self.embedding nn.Embedding(config[vocab_size], config[d_model]) self.pos_encoding PositionalEncoding(config[d_model], config[max_len], config[dropout]) self.layers nn.ModuleList([ TransformerBlock(config) for _ in range(config[n_layers]) ]) self.norm nn.LayerNorm(config[d_model]) def forward(self, src, src_maskNone): x self.embedding(src) * math.sqrt(config[d_model]) x self.pos_encoding(x) for layer in self.layers: x layer(x, src_mask) return self.norm(x)这里有一个细节self.embedding(src) * math.sqrt(d_model)。这一步在论文里写得很明白但很多代码注释都会略过。因为 embedding 是随机初始化的它的方差比较小乘上sqrt(d_model)是为了让 embedding 的数值范围和位置编码匹配防止位置编码把信号淹掉。这属于“不加也行但加了更稳”的工程细节。src_mask通常是 padding mask因为一个 batch 里的句子长度不一样短句子后面要用padtoken 补齐而这些 padding 位置不应该参与注意力。实现上mask 的形状一般是[batch, 1, 1, src_len]它会沿着注意力分数的形状做广播把 padding 位置全部屏蔽掉。3.2 TransformerBlock 的完整实现与形状追踪为了让打印形状变得更方便我在TransformerBlock里会故意保留几个临时变量方便在 notebook 里观察数据变换class TransformerBlock(nn.Module): def __init__(self, config): super().__init__() self.attn MultiHeadAttention(config[d_model], config[n_heads], config[dropout]) self.ffn FeedForward(config[d_model], config[d_ff], config[dropout]) self.norm1 nn.LayerNorm(config[d_model]) self.norm2 nn.LayerNorm(config[d_model]) self.dropout nn.Dropout(config[dropout]) def forward(self, x, maskNone): # 第一段多头注意力 残差 attn_out, _ self.attn(x, x, x, mask) x x self.dropout(attn_out) x self.norm1(x) # Pre-Norm先归一化 # 第二段前馈网络 残差 ffn_out self.ffn(self.norm2(x)) x x self.dropout(ffn_out) return x看起来顺序有些别扭常见的架构图是先画 AddNorm再进 FFN这里我在 FFN 之前又加了一次self.norm2(x)。这其实就是 Pre-Norm 的标准写法每个子层输入前先归一化残差连接在外面。你可以把它改造成 Post-Norm也就是归一化放到残差后面。两种都能跑但你会发现 Pre-Norm 对学习率更宽容在层数更深时表现更稳定。跑一次前向你会看到维度全程不变输入是[batch, seq_len, d_model]中间每一层的输出仍然是[batch, seq_len, d_model]。这就给了你一个巨大的调试便利无论模型中间哪一段出问题你只需要比较它的输入输出维度是否保持一致。我实际带过的一些朋友最初写代码时反复在 shape 上报错后来养成了“每层打印 shape”的习惯五分钟内就能定位问题。3.3 Decoder 的差异因果掩码与 Cross-AttentionEncoder 和 Decoder 的代码结构高度相似区别主要在两个地方。第一个是 Decoder 做的是“自回归”生成预测第 i 个 token 时只能看前 i-1 个 token不能偷看未来。这个限制靠因果掩码causal mask实现代码通常是causal_mask torch.tril(torch.ones(max_len, max_len)).view(1, 1, max_len, max_len).bool()torch.tril生成一个下三角矩阵矩阵里位置[i, j]为 1 表示第 i 个 token 能看到第 j 个 token当 j i 时。在注意力分数矩阵里这个 mask 会把它右上角的未来信息全部变成-inf。调试时最容易出的问题是 mask 维度没对齐scores的形状是[batch, heads, tgt_len, src_len]所以 mask 至少要能广播成这个形状建议维度是[batch, 1, tgt_len, src_len]或[1, 1, tgt_len, src_len]。第二个区别是 Decoder 里的注意力有两种第一层是自回归自注意力看的是自己之前生成的 token第二层是 Cross-AttentionQ来自 DecoderK和V来自 Encoder 的输出。这就是“翻译”的关键Decoder 从目标语言里拿到已经生成的部分同时从源语言里拿到已经编码好的完整语义。代码里Cross-Attention 其实不需要额外写一个新类直接复用MultiHeadAttention只要 forward 时传入不同的 K 和 V 就行enc_out encoder(src) dec_out decoder(tgt, enc_out, tgt_mask, src_mask)我个人做调试时会刻意把 Encoder 和 Decoder 的输出分开打印。因为在实际训练初期Decoder 的 loss 下降速度通常比 Encoder 慢得多如果一开始就合在一起训练你会很难判断是编码端没学好还是解码端没学好。虽然很多人直接端到端训练但对刚接触代码的人来说单独跑一遍两个模块能让你理解得更透。4. 训练与推理实战一个小模型的完整生命周期4.1 数据准备与 padding mask 的生成训练一个 mini Transformer我建议用一个小型平行语料比如英法简表、或者随便一对句子序列重点不是效果多好而是让你完整跑通流程。原始文本要先分词、转成 id、按 batch 拼装成定长张量。这里有一个经典问题每个 batch 是最长句子长度还是统一固定长度实践中我一般用动态 padding也就是每一个 batch 只 pad 到该 batch 内最长句子的长度。这样不会浪费大量算力在padtoken 上。生成 padding mask 的代码也很简单def make_pad_mask(seq, pad_idx0): # seq: [batch, seq_len] return (seq ! pad_idx).unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len]注意这个 mask 最终要变形到注意力分数能广播的形状。如果你的scores是[batch, heads, tgt_len, src_len]那[batch, 1, 1, src_len]就可以沿着heads和tgt_len广播。4.2 组合训练循环与损失计算训练循环写出来后你会发现 Transformer 和别的神经网络没有本质区别前向、算 loss、反向、更新。但有两个细节要专门说。第一交叉熵损失计算时nn.CrossEntropyLoss需要直接把 logits 扔进去它内部会做 softmax如果你在外面又做一次 softmax反而会算错梯度。第二padding 位置不能被算进 loss所以要把损失函数的ignore_index设为pad的 id。criterion nn.CrossEntropyLoss(ignore_index0) optimizer torch.optim.Adam(model.parameters(), lr1e-3, betas(0.9, 0.98), eps1e-9)betas(0.9, 0.98)是论文里推荐的超参它和我们平时常用的(0.9, 0.999)略有不同。更稳妥的做法是给 Adam 配一个学习率预热warmup前 4000 步让学习率从很小的值线性涨到峰值之后再按步数衰减。因为 Transformer 的初始化方差相对较大如果一开始就用大学习率容易出现早期 loss 震荡甚至直接发散。训练几个 epoch 后如果 loss 在稳定下降你会看到它从最初的 10 左右一路降到 2 以内。遇到 loss 突然变成 NaN 时先别急着调模型把学习率调低一个数量级再试。我见过大量 NaN 案例九成是学习率太高或注意力没有缩放只有一成是数据里有无穷值。4.3 推理阶段与贪心解码训练完毕推理阶段就是自回归式地一个词一个词往外蹦。核心逻辑是给定起始 tokenbos喂进 Decoder 得到下一个 token 的概率取概率最高的那个拼到序列后面再重复这个过程直到生成eos或者达到最大长度。伪代码如下def greedy_decode(model, encoder_output, src_mask, start_token, end_token, max_len50): tgt torch.tensor([[start_token]]).to(src_mask.device) for _ in range(max_len): tgt_mask make_causal_mask(tgt.size(1)).to(src_mask.device) logits model.decode(encoder_output, tgt, tgt_mask, src_mask) # [1, tgt_len, vocab] next_token logits[:, -1, :].argmax(dim-1) if next_token.item() end_token: break tgt torch.cat([tgt, next_token.unsqueeze(0)], dim-1) return tgt有一个很常见的 bug生成时忘记更新tgt_mask的长度。tgt每增加一个 token因果 mask 的尺寸也要跟着变否则新位置看不到旧位置或者根本找不到它的 mask直接报维度错误。这个问题我在自己写的第一个翻译 demo 里也踩过后来索性把 mask 生成函数做得更“自动化”一点直接根据tgt.size(1)动态生成省得每次手滑。4.4 训练日志怎么读loss、困惑度与过拟合信号看训练日志不能只看 loss还要结合困惑度perplexity来判断模型是不是真的在学。困惑度等于exp(loss)。如果 loss 到 2.0困惑度就是 7.4意思是模型平均在每个位置上对 7.4 个候选词感到困惑这算一个中等水平如果 loss 到 1.0困惑度约 2.7说明模型已经比较有把握了。对一个小型 demo 来说不要指望它达到大模型那种 1.0 以下的 loss更重要的指标是验证集 loss 是否和训练集 loss 同步下降。一旦发现训练 loss 还在降、验证 loss 开始回升恭喜你过拟合了。这时候应该加大 dropout、减小模型容量或增加数据量。5. 常见坑位排查注意力 NaN、mask 错位与头数不匹配5.1 shape 不匹配与维度混淆速查我把常见的 shape 错误整理成了一张速查表希望你下次遇到报错能直接对号入座报错信息常见原因解决办法Expected 4D input注意力打分时 Q 和 K 没有扩到 4 维确认view后已transpose(1,2)The size of tensor a (8) must match size of tensor b (16)n_heads与d_model不整除将d_model除以n_heads时确保是整数mask shape mismatchmask 维度不能广播到scores统一 mask 为[batch, 1, tgt_len, src_len]CUDA error: device-side assert triggered通常是ignore_index没设好或 token id 越界检查 label 里是否混入非法 idone element cannot be broadcastedEmbedding 与位置编码 seq 维度不匹配切取pe[:, :x.size(1)]头数不匹配这个坑我在很多小白代码里反复看到。d_model512n_heads6这时候 512 不能被 6 整除得到的d_k是 85.33你就没法用view拆分了。要么把n_heads改成 8要么把d_model改成 480 或 504。工程上很少为了迁就头数去改d_model基本都是把头数调到能整除的数值。5.2 注意力 NaN 的定位方法NaN 是最让人头大的问题因为它不会告诉你具体是哪一行。我的排查顺序是先看 loss 在哪一步变成 NaN再在 forward 里插入中间打印检查 attention scores 和 softmax 输出的数值范围。如果是scores爆炸检查有没有除以sqrt(d_k)如果是scores里出现-inf检查 mask 是不是把该看的 token 也屏蔽了。还有一个容易被忽略的源头FFN 里的 ReLU 输出范围太大。如果d_ff很大ReLU 之后有些位置的数值可能是几十甚至上百再加上残差连接数值一路滚雪球。这也是很多现代实现把 ReLU 换成 GELU 或 SwiGLU 的原因之一。GELU 在负半轴更平滑不会像 ReLU 那样直接一刀切。如果训练中早期 loss 一直上蹿下跳可以考虑把激活函数换成nn.GELU()。5.3 mask 错位的三类典型表现mask 错位通常有三种表现生成时出现重复词、训练 loss 比预期高、或者模型在测试时“偷看未来”导致生成的连贯性很差。我遇到最隐蔽的一种是mask 形状对了但数值反了。比如本该屏蔽的位置填了 1该保留的位置反而被屏蔽成了 0。判断技巧很简单把 mask 打印出来确认下三角矩阵的 1 在左下角而右上方是 0。拿tgt_len4来说合法 mask 应该是[[1, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]这个矩阵的生成方式有很多种但核心只有一条i行j列表示第i个输出 token 是否能看到第j个输入 token。只要这个逻辑搞对了mask 基本不会错。5.4 残差连接与 LayerNorm 的位置混淆Pre-Norm 和 Post-Norm 的实现差异虽然只有几行代码但对训练行为影响很大。Post-Norm 的写法更贴近原始论文x Norm(x Sublayer(x))。Pre-Norm 的写法则是x x Sublayer(Norm(x))。如果同时使用大学习率和 Post-Norm浅层容易出现梯度不稳定Pre-Norm 则相对友好。但 Pre-Norm 有一个副作用最后一层输出和残差直接相加导致输出数值可能没有被充分归一化。所以很多实现会在最后加一个Final LayerNorm。如果你搭建的模型最后一层后直接接nn.Linear到词表我的建议是加一层nn.LayerNorm(d_model)作为收尾能显著提升稳定性。这也是我在看很多开源代码时发现的一个共性设计不是论文里写的但确实是大家踩过坑后总结出的经验。6. 从代码到理解注意力可视化、参数量与扩展方向6.1 用一张注意力热力图验证模型真的在“学”模型跑起来以后很多人就想看一眼 attention 到底学到了什么。最简单的方法是让MultiHeadAttention的 forward 返回权重矩阵然后画热力图。你会看到不同头有不同的关注模式有些头对角分布说明它主要盯紧当前位置和前一个位置有些头有较宽的横向条纹说明它在做长距离的实体对齐。一个小技巧是把注意力权重和对应的句词打印出来对照着看。我曾经在一个中文翻译 demo 里发现某个头特别关注“你”和“我”这类代词另一个头则总是把注意力均匀撒在标点和停用词上。这种可视化并不仅仅是为了“好看”它能帮你确认训练是否正常。如果一个头在所有层里的注意力权重都是均匀分布可能性有两种要么还没收敛要么这个头已经“死”掉了。6.2 参数量估算为什么 512 维这么流行用手算一遍参数量能让你对模型规模有更具体的感知。Embedding 的参数量是vocab_size * d_model1 万词表加 512 维就是 512 万。多头注意力的三个投影加一个输出投影每个都是512*512总共约 104 万。FFN 的两个线性层则是512*2048 2048*512约 209 万。单层 Block 大约 300 多万参数6 层叠加后主体就是 2000 多万再加上 Embedding 与其他一个小型 Transformer 大概在 2500 万参数左右。这个规模在普通显卡上完全能训起来但如果你想带着 3 万词表跑光 Embedding 就快 1.5 亿参数内存压力就上来了。现在你也能理解为什么d_model512是经典之选它既能表达足够丰富的信息又让多头注意力和 FFN 的参数量维持在可控范围内。真的跑到 GPT 级别的模型d_model往往在 768 到 12288 之间参数量随维度平方级膨胀这也是为什么大模型动辄几百亿参数。6.3 为了让代码更现代Flash Attention、相对位置编码和 KV Cache如果你已经跑通基础版本我建议你接着研究三个扩展点。第一个是 Flash Attention它的核心不是改变注意力公式而是通过分块计算和重计算减少显存占用、提高速度。原来的注意力需要先把[seq_len, seq_len]的完整分数矩阵算出来并存到显存里Flash Attention 则把这块矩阵切成小块边算边用内存占用从平方级降为线性级。第二个是相对位置编码经典的正弦编码已经不太够用。很多新模型用旋转位置编码RoPE或 ALiBi它们不仅编码绝对位置还能让模型更好感知 token 之间的相对距离。第三个是 KV Cache在推理阶段把已经算过的K和V缓存起来避免每生成一个新 token 就把过去所有位置的K、V重算一遍。推理性能的提升非常明显尤其是长文本生成时几乎是从“不能忍”变成“丝滑流畅”。6.4 我的个人实践心得讲真Transformer 代码本身并不难难的是你容易被各种花哨的工程优化吓到。我见过太多人花了一个星期去追最新的稀疏注意力实现结果最后连最基础的scaled_dot_product_attention都写不顺。我的体会是先把最朴素的版本跑通再一步步往上加工程技巧这样你才能分清哪些修改是真正必要的哪些只是锦上添花。另外我很推荐你在调试时顺手加一个“形状断言工具”——可以是你自己的函数也可以是 PyTorch 内置的assert。每次输入输出经过某个子层时都检查维度是否符合预期。这个过程很笨但极其有效。我在实际项目里用它节省的时间远比写它用的时间多得多。一个小技巧是给每个子层起一个简短的名字比如emb、pos、attn、ffn然后打印形状时带着名字输出这样出错时一眼就能看见是哪段炸了。最后再分享一个小经验如果你真想彻底搞懂 Transformer不要只盯着别人的代码看挑一个周末自己从头写一遍写完再跑一遍机器翻译的 demo。这个过程的收获比你看十篇解析文章都大。写不出来也没关系卡住的地方恰恰是你理解最薄弱的地方把它们标记下来再回头看我前面讲的这些顺序去补课你会发现自己突然就能看懂那些开源巨作的源码了。