之前有段时间我觉得自己把Transformer相关的东西学得差不多了QKV、自注意力、位置编码一套一套都能讲。结果某天想读一份开源实现打开文件没看几行就被各种 reshape、transpose 和 mask 绕晕了。后来我换了个思路——不去硬啃别人的完整源码而是自己动手写一个最简版本。这篇文章就是我用深度学习框架手写一个“能跑版” Transformer 的完整记录代码量不大但足够把核心模块串起来。如果你和我一样属于“论文看懂不少、代码没写过几行”的状态建议跟着把代码敲一遍比看十篇解析都管用。1. 先从整体结构下手拿到一份Transformer代码先看什么1.1 先明确这个Transformer是哪种TransformerTransformer 在原始论文里是 Encoder-Decoder 结构但后来大家在代码里见到了更多变体Encoder-only比如做分类、Decoder-only比如做文本生成。我这次实现的是一个 Decoder-only 的简化版本因为它包含了 Transformer 最核心的几个组件多头自注意力、位置编码、前馈网络、LayerNorm、残差连接。你把这套代码看懂以后想改成 Encoder-only 其实非常简单——把因果 mask 拿掉让每个位置看到全序列就行想加 Cross-Attention也只需要在 TransformerBlock 里再插一个注意力模块。我对“解析代码”这件事的标准是必须让我能在笔记本上把模型跑起来而不是只贴一个永远没办法运行的半成品。所以下面所有代码都是完整可复现的你把它按顺序拼到一个文件里就能训练。模型很小但“麻雀虽小五脏俱全”核心链路一个都不少。1.2 整条数据流和shape变化先别急着看代码闭眼想一遍数据是怎么流的。输入是一个整数序列假设形状是 (B, T)通过 token embedding 查表变成 (B, T, C)加上位置编码还是 (B, T, C)经过 N 个 Transformer Block每个 Block 内部都是“注意力 残差 LayerNorm 前馈 残差”输出形状不变最后过一个线性层变成 (B, T, vocab_size)得到每个位置预测下一个 token 的 logits算交叉熵损失。所有 shape 变化我建议用一张表记下来后面调试代码时对照着看阶段张量形状说明输入 tokens(B, T)整数 idtoken embedding(B, T, C)C 是 d_model加位置编码(B, T, C)每个位置一个向量注意力分数(B, n_heads, T, T)每两个 token 之间的相关度注意力输出(B, T, C)多头合并后前馈输出(B, T, C)逐 token 变换最终 logits(B, T, vocab_size)预测分布代码里我习惯用 B、T、C 三个字母表示 batch size、序列长度、模型维度。C 就是 d_model。很多开源实现也这么命名你抓住这套维度关系后面任何 Transformer 代码都能拆开。2. 把多头注意力拆成代码别再背诵QKV公式2.1 为什么需要多头而不是一个头在动手写代码之前先建立一点直觉。单头注意力就是一组 Q、K、V它只能学一种“谁和谁有关”的模式。多头相当于把模型分成多个子空间让不同的头分别关注不同类型的关系有的头关注相邻词有的头关注跨长距离的指代有的头关注某种句法结构。最后再把所有头的输出拼起来让模型拥有更丰富的表达能力。从公式上看多头注意力就是缩放点积注意力Attention(Q, K, V) softmax(QK^T / sqrt(d_head)) V差别在于Q、K、V 先被拆成 n_heads 份每份维度是 d_head d_model / n_heads各部分独立计算注意力最后把结果 concat 回去。这也是为什么代码里会有大量 view 和 transpose。2.2 MultiHeadAttention类的实现与逐行解释直接看代码。这个类大概是我们整个 Transformer 里最核心的部分我建议你亲手敲一遍import torch import torch.nn as nn import torch.nn.functional as F 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.head_dim d_model // n_heads self.wq nn.Linear(d_model, d_model, biasFalse) self.wk nn.Linear(d_model, d_model, biasFalse) self.wv nn.Linear(d_model, d_model, biasFalse) self.out_proj nn.Linear(d_model, d_model, biasFalse) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): B, T, C x.shape # 线性投影并拆成多头 q self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) # 缩放点积注意力 attn_scores q k.transpose(-2, -1) / (self.head_dim ** 0.5) if mask is not None: attn_scores attn_scores.masked_fill(mask 0, float(-inf)) attn_weights F.softmax(attn_scores, dim-1) attn_weights self.dropout(attn_weights) # 加权聚合然后把多头合并 out attn_weights v out out.transpose(1, 2).contiguous().view(B, T, C) # 输出投影 return self.out_proj(out)有几个点必须拆开讲。第一为什么线性投影不加 bias有些实现会加但后面有 LayerNorm 和残差连接bias 在多数情况下不是必须的。这里我习惯关闭减少一点参数也不会影响收敛。第二view 和 transpose 的顺序。x 从 (B, T, C) 变成 (B, T, n_heads, head_dim)这一步只是拆成 4 维但此时 head 还在第 2 维我们需要把它换到第 1 维让每个 head 独立计算所以 transpose(1, 2)。最终变成 (B, n_heads, T, head_dim)。所有头在 batch 维之后并行计算这正是多头注意力的精髓。第三缩放因子是 head_dim ** 0.5不是 d_model ** 0.5。这一点容易被忽略。为什么要缩放因为 QK^T 的数值会随着维度和长度变大而变大不缩放的话 softmax 会进入饱和区梯度变得很小。d_head 决定了每个头上点积的尺度所以用 head_dim 来缩放。2.3 因果mask让Decoder承担自回归任务的关键Decoder 在生成每个 token 时绝对不能看到未来的 token否则就“作弊”了。训练的时候虽然我们可以一次性把整条序列喂进去但通过一个上三角全为 0 的 mask让注意力分数在未来的位置变成负无穷softmax 之后这些位置的概率就变成 0。因果 mask 的构造非常简单def create_causal_mask(T): mask torch.tril(torch.ones(T, T)).bool() return maskmask 的形状是 (T, T)其中 mask[i][j] True 表示位置 i 可以 attend 到位置 jFalse 表示不能。这个 mask 在传入前向函数后会自动广播到 (B, n_heads, T, T) 上不需要额外扩展。一个常见的疑问是为什么训练时不能像推理那样一个个 token 循环生成因为训练需要高效并行。Transformer 本身是定长序列到定长序列的计算mask 把它伪装成“因果”的但内部所有位置的注意力计算还是同时进行的。这也是自回归语言模型训练效率高的原因。3. 位置编码代码里最容易忽略却最讲究的一层3.1 经典sin-cos位置编码的代码落地Transformer 没有循环结构也没有卷积如果直接把 token embedding 输入注意力模型完全分不清“我喜欢你”和“你喜欢我”在位置上的差别。位置编码就是为了给每个位置一个独特的“指纹”。经典正弦位置编码的公式是PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))代码实现如下import math def sinusoidal_position_encoding(max_len, d_model): assert d_model % 2 0 pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2, dtypetorch.float) * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe代码中有个细节计算10000^(2i/d_model)时直接调pow容易出现数值问题尤其是 i 很大的时候。先取对数再取指数也就是exp(-log(10000)*2i/d_model)在数值上更稳定。div_term的 shape 是 d_model/2分别作用在偶数维和奇数维上。实际加到输入上时只需要截取当前序列长度x x self.position_encoding[:T]在注册为 buffer 的情况下position_encoding会跟着模型一起切换设备不需要手工.to(device)。3.2 为什么不能直接用数字1、2、3最容易想到的位置编码是给每个位置加一个标量位置 0 加 0位置 1 加 1位置 2 加 2。这个方案的主要问题有两个。第一无界。位置编码的值会随序列长度无限增长可能掩盖 token 本身的语义信息。一个数值很大的位置向量会“淹没” embedding。第二模型很难从无界整数中泛化到更长的序列。sin-cos 编码把位置映射到一个有界区间内并且相邻位置之间保持平滑的相似性模型更容易学习相对位置的模式。事实上sin-cos 编码还具备一些漂亮的线性性质通过变换可以表达相对位置关系这是直接加序号做不到的。3.3 可学习位置编码和正弦位置编码怎么选在实际代码里你还会看到另一种做法把位置编码当作一个nn.Embedding随机初始化然后跟着模型训练。self.position_embedding nn.Embedding(max_len, d_model)两种方案各有取舍。正弦编码完全不需要额外参数有更好的外推能力可学习编码在训练数据足够多时表现也不错且实现更简单。近年很多大模型直接用可学习位置编码因为它们的训练语料极其充足外推需求不是首要考虑。但如果你是解析经典 Transformer 代码肯定会遇到 sin-cos 版本所以这一节务必要吃透。4. 残差、LayerNorm和前馈网络Transformer块的骨架4.1 Pre-LN还是Post-LN一个顺序决定训练难度原始论文里的顺序是先做注意力再加残差最后 LayerNorm也就是Post-LNx LayerNorm(x Attention(x))后来大量实践发现Post-LN 在深层模型上训练不稳定往往需要学习率预热warmup来稳定训练。现代实现更常见的是 Pre-LNPre-LNx x Attention(LayerNorm(x))Pre-LN 把归一化放在前面残差路径始终保持“干净”的恒等映射梯度可以直接从最后一层回流到第一层训练稳定很多也不需要复杂的预热策略。代码上只是几行的顺序差异但影响非常大。我们这里的示例使用 Pre-LN这也是我现在写模块的默认选择。4.2 FFN的d_ff为什么要膨胀多头注意力本质上是一个加权求和权重来自输入之间的相似度这个过程是线性的。如果堆叠很多层注意力模型整体表达能力还是受限。前馈网络的作用就是给每个 token 做一次非线性特征变换通常的做法是先把 d_model 映射到一个更大的中间维度 d_ff激活后再映射回来。为什么中间维度要变大可以把它理解为一个“展开再压缩”的过程。展开阶段把每个位置的局部特征扩展成更高维的表达激活函数引入非线性压缩阶段再把高维特征提炼成下一层需要的语义信息。经典配置是 d_ff 4 * d_model这个比例在绝大多数代码里都成立。实现也很简洁class FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.net nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), nn.Dropout(dropout) ) def forward(self, x): return self.net(x)原始论文用的是 ReLU现在很多实现喜欢 GELU。如果看到老代码里的 ReLU也不要惊讶功能上都是非线性激活只是曲线平滑度不同。4.3 把模块拼成TransformerBlock有了注意力、LayerNorm、前馈网络一个 Transformer Block 就是把他们按 Pre-LN 的顺序粘起来class TransformerBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn FeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) def forward(self, x, maskNone): x x self.attn(self.norm1(x), maskmask) x x self.ffn(self.norm2(x)) return x这里的 LayerNorm 只对最后一维 C 做归一化不会跨 batch 或序列长度。这个细节很重要因为很多刚接触的人会把 LayerNorm 和 BatchNorm 搞混。LayerNorm 是对每个 token 的向量内部做归一化和序列中其他 token 无关。残差连接的价值也值得强调。没有残差连接几十层堆叠之后梯度很容易消失有了残差每一层都保证输入信息能直接向后传递。你可以把残差理解成一个“高速公路”模型可以选择只在这个高速公路上加一点微小的变换也可以利用它把信息传到很深的层。5. 从零拼出一个能跑的Transformer训练小例子5.1 设计一个能验证模型是否真在学习的任务理论学习到最后最重要的就是模型能跑起来。为了快速验证代码的正确性我找了一个简单到不能再简单的任务让模型学习循环序列。具体来说我们用数字 0 到 4 不断重复0, 1, 2, 3, 4, 0, 1, 2, 3, 4, ...然后给模型一段长度为 T 的窗口让它预测下一个 token。因为规律非常明显只要模型真的学会了位置编码和注意力loss 应该快速下降。这个任务不需要大模型也不需要大规模训练数据非常适合验证 Transformer 代码链路。数据生成代码不需要写 DataLoader 类几行函数就够了def random_batch(seq, batch_size, T): starts torch.randint(0, len(seq) - T - 1, (batch_size,)) xs, ys [], [] for s in starts: xs.append(seq[s:s T - 1]) ys.append(seq[s 1:s T]) return torch.stack(xs), torch.stack(ys)5.2 模型结构参数怎么定因为任务很简单我用非常小的模型。选择参数要遵守一个原则尽量小到 CPU 能跑但大到足够体现多头和深度的作用。下面这组参数是我试过比较舒服的配置vocab_size 5只有 0 到 4 五个数字d_model 32n_heads 4head_dim 8d_ff 64num_layers 2block_size 64最大序列长度注意 d_model 必须能被 n_heads 整除否则注意力模块的 assert 会直接报错。d_model 选 32 而 head_dim 是 8这个维度对 CPU 上的小任务很友好同时还能展示多头并行。完整模型类class Transformer(nn.Module): def __init__(self, vocab_size, d_model, n_heads, d_ff, num_layers, max_len64, dropout0.1): super().__init__() self.token_embedding nn.Embedding(vocab_size, d_model) self.position_encoding sinusoidal_position_encoding(max_len, d_model) self.register_buffer(pe, self.position_encoding) self.blocks nn.ModuleList([ TransformerBlock(d_model, n_heads, d_ff, dropout) for _ in range(num_layers) ]) self.norm_final nn.LayerNorm(d_model) self.lm_head nn.Linear(d_model, vocab_size) def forward(self, tokens): B, T tokens.shape x self.token_embedding(tokens) x x self.pe[:T] mask torch.tril(torch.ones(T, T, dtypetorch.bool, devicetokens.device)) for block in self.blocks: x block(x, mask) x self.norm_final(x) logits self.lm_head(x) return logits位置编码在__init__里通过register_buffer注册这样模型调用.to(device)或.cuda()时位置编码也会自动跟着移动省去手动转移设备的麻烦。5.3 训练循环和loss观察训练循环非常常规。优化器用 Adam学习率设 1e-3损失函数是交叉熵model Transformer(vocab_size5, d_model32, n_heads4, d_ff64, num_layers2, max_len64) opt torch.optim.Adam(model.parameters(), lr1e-3) crit nn.CrossEntropyLoss() pattern torch.arange(5).repeat(100) for step in range(500): x, y random_batch(pattern, batch_size32, T8) opt.zero_grad() logits model(x) loss crit(logits.transpose(1, 2), y) loss.backward() opt.step() if step % 100 0: print(fstep {step}: loss {loss.item():.4f})运行之后你会发现loss 一开始大约在 1.6 附近因为 5 个类别的均匀分布熵就是 log5≈1.609。这意味着模型初期完全没有学习。几百步内 loss 会明显下降最后降到很低说明模型已经掌握了循环规律。如果 loss 一直不下降别怀疑任务先回头查上一节的那些 shape 和 mask 细节。5.4 生成让模型自己预测下一个token训练完之后可以写一个最简单的贪心生成函数看看模型能不能自己续写def generate(model, start_tokens, max_new20, block_size64): model.eval() tokens list(start_tokens) with torch.no_grad(): for _ in range(max_new): inp torch.tensor([tokens[-block_size:]], dtypetorch.long) logits model(inp) next_id logits[0, -1].argmax().item() tokens.append(next_id) return tokens这段代码每次把当前已有的 token 序列重新送进模型只取最后一个位置的预测结果然后拼回去。这个“重新前向计算整个序列”的做法也是很多生成模型推理速度慢的原因因为每生成一个 token 都要全量计算一遍。gpt 类模型生成新词时能缓存历史的 K/V 从而加速但这是优化层面的问题不影响理解。6. 调试Transformer代码时踩过的坑shape、mask、NaN与顺序6.1 shape不匹配先画表再写代码我写这份代码时犯过最多的错误就是 shape 问题。最常见的是view和transpose的顺序搞反。很多新手会把x.view(B, n_heads, T, head_dim)直接写出来但这是错的因为原来的最后一个维度是 C要先把 C 拆成 (n_heads, head_dim)然后再把 n_heads 换到第 1 维。我的调试习惯是在关键位置打印 shape。比如注意力模块里打印q.shape、k.shape、attn_scores.shape和out.shape一行行对表。写代码不是一次就能跑通的打印调试是最快的路径。6.2 mask的广播一点要小心mask 的形状是 (T, T)但注意力分数是 (B, n_heads, T, T)。很多新手会觉得必须手动把 mask 扩展成 (B, n_heads, T, T) 才能用实际上masked_fill会自动广播不需要额外处理。真正容易翻车的是 mask 的逻辑方向。我习惯用 bool 类型的 maskTrue 表示可见、False 表示不可见。这样在代码里写masked_fill(mask 0, float(-inf))非常直观。如果你用的是 float mask0 和 1 的含义就要格外小心很容易弄反。还有一个隐蔽问题如果某个位置的所有注意力分数都被 mask 成-infsoftmax 之后可能出现 NaN。正常情况下因果 mask 不会出现整行都被 mask 的情况但如果你在建 mask 时方向写反就很有可能。6.3 loss不降和NaN的排查顺序遇到 loss 不降我一般按这个顺序排查检查 label 是否右移正确。自回归模型里输入是前 T-1 个 token目标必须是后 T-1 个 token。检查 mask 是否让每个位置都能至少看到一个 token。检查多头合并时有没有用 contiguous 再 view。把(B, n_heads, T, head_dim)转回(B, T, C)时先transpose再view会报错必须先.contiguous()。检查激活函数和初始化。小模型一般不需要特殊初始化但如果 loss 变成 NaN最常见原因是学习率太大。加一行torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)能避免大部分梯度爆炸。6.4 结构性建议先跑最小case再放大最后分享一个我自己一直在用的调试技巧先跑 batch_size1、序列长度1确认 forward 能过再跑 batch_size2、序列长度3确认多头和 mask 没有逻辑错误最后才用完整配置训练。这样做可以把问题从“模型结构错误”和“训练参数问题”中快速分开。还有一个小经验写 Transformer 模块时不要把 LayerNorm 放在forward函数里临时创建而要在__init__里注册为子模块。否则模型的参数无法被优化器管理训练时会出现“模型 loss 下降但参数不更新”的诡异现象。听起来很蠢但我真的见过有人把 norm 写成每次 forward 都新实例化然后怎么调都不收敛。代码这个东西跑通了并不难难的是把每个“为什么”都想明白。如果你按这篇文章把代码敲下来并且真的让 loss 下降了我相信你再去读那些大而全的开源实现至少不会一看到 reshape 就发怵。写 Transformer 代码没有捷径能做的就是把每一步 shape 和每一块 mask 都亲手走一遍。