手撕 Transformer最值得先看的不是把整份代码背下来而是搞清楚输入、注意力、编解码器这三条主线怎么串起来。这篇文章用 PyTorch 从零实现一遍输入经过 embedding、位置编码变成向量自注意力在向量之间交换信息编码器输出上下文表示解码器再靠着因果掩码和交叉注意力逐步生成结果。适合已经会用 PyTorch 搭分类网络、但对 Transformer 内部维度还比较模糊的读者如果你只是调用 nn.Transformer 一次跑通那不算手撕真正自己把 MultiHeadAttention 和 mask 写一遍才能理解维度为什么是这个形状、训练时为什么需要 teacher forcing。我会尽量把每一段代码都控制在“能跑、能看、能改”的程度。你不需要一次性把全部代码背下来先按顺序跑通再回头改参数效果会比反复看原理图更明显。1. 先拆清三条主线输入、注意力、编解码器各负责什么1.1 为什么建议手撕而不是直接调 nn.Transformer很多人学 Transformer 时会有一种错觉用官方封装跑一次翻译任务就以为理解了。实际上模型能跑通不等于你理解了自己写的代码。nn.Transformer把输入映射、多头注意力、mask、LayerNorm 全封装在内部你只能看到输入输出形状看不到 Q/K/V 在中间发生了什么。自己手撕等于把这些问题暴露出来输入 token 怎么从整数变成向量多头注意力里的“多头”是怎么拆开又是怎么合并回去的编码器层和解码器层之间哪些张量是共享哪些是每次重新计算的解码器训练时为什么要用掩码推理时为什么要一个 token 一个 token 地生成这些问题在调包时不会出现。但只要训练 loss 不降、显存溢出、生成结果乱码你迟早要回来看内部实现。1.2 这篇文章的整体验证思路我的建议是不要一开始就追求“完整复现论文里的 base 模型”。先把三个模块按顺序搭出来输入映射层Embedding PositionalEncoding。注意力层Scaled Dot-Product Attention Multi-Head Attention。编解码器层EncoderLayer、DecoderLayer、Transformer 主类。每写完一个模块就打印输入输出形状。形状对了再拼下一个模块。等整体拼接完成用一个很小的玩具任务验证训练流程比如让模型把一串数字逐个加一。这个任务本身没有实际价值但它能快速暴露 mask、维度、loss 计算这一类基础问题。注意第一次测试尽量用 CPU 跑batch 设成 2序列长度设成 8d_model 设成 32。先看形状和流程不要在一开始就纠结训练效果。2. 环境准备与验证思路为什么先跑通最小配置2.1 依赖和版本怎么确认手撕 Transformer 不需要太重的基础环境。能用 PyTorch 建一个 nn.Module能执行 tensor 操作就足够了。建议 Python 3.9 以上PyTorch 用 2.x 或你本地已有的稳定版本都可以核心用到的都是很基础的接口nn.Embeddingnn.Linearnn.LayerNormnn.Dropouttorch.matmul、F.softmax不强制要求 GPU。模型规模在几十万参数以下时CPU 跑训练也很快。安装 PyTorch 时如果不想折腾 CUDA可以先装 CPU 版本后面需要 GPU 再换对应版本。重要的是先确认 PyTorch 安装成功import torch print(torch.__version__)能正常打印版本号就说明环境可用。接下来用一行代码验证基本操作x torch.randn(2, 8, 32) print(x.transpose(1, 2).size())这里2是 batch8是序列长度32是特征维度。后面写代码时所有报错几乎都围绕这三个数字展开。2.2 第一次验证的任务规模我建议采用这样一组最小配置d_model 32n_heads 4head_dim d_model // n_heads 8num_encoder_layers 2num_decoder_layers 2d_ff 64这组配置不是用来刷效果的而是为了让每一步张量形状都容易打印、容易检查。如果要跑到论文里常见的 d_model512、n_heads8代码逻辑一样只是显存和训练时间明显上升。低配置能跑通不代表适合批量训练。模型越大学习率、dropout、归一化位置的作用就越敏感。3. 输入处理Embedding、位置编码和 batch 张量形状3.1 从 token 到 embedding为什么还要乘 sqrt(d_model)Transformer 的输入不是普通矩阵而是 token 序列。假设词表大小为 vocab_size输入是一个 batch 的整数序列shape 是(batch_size, seq_len)。第一步就交给nn.Embeddingself.embedding nn.Embedding(vocab_size, d_model)这样每个 token 变成一个d_model维向量输出 shape 变为(batch_size, seq_len, d_model)。在原始论文里embedding 之后还会乘上math.sqrt(self.d_model)。为什么要这么做因为后面马上要加位置编码embedding 的数值如果太小位置信息会盖过 token 本身的信息。缩放一下可以让嵌入向量和位置编码在量级上更接近后续注意力计算也更稳定。这不是必须的但复现论文实现时通常会保留。3.2 位置编码用 sin/cos 还是可学习参数Transformer 不像 RNN 那样天然知道 token 顺序。自注意力对位置不敏感它只看 token 两两之间的相似度。所以要把位置信息显式加进输入。经典做法是用 sin/cos 的固定位置编码import math import torch import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000): super().__init__() pe torch.zeros(max_len, d_model) position torch.arange(0, max_len).unsqueeze(1).float() div_term torch.exp( torch.arange(0, d_model, 2).float() * (-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) self.register_buffer(pe, pe) def forward(self, x): return x self.pe[:, : x.size(1)]这里把pe注册成 buffer意思是它不会参与梯度更新但会跟着模型一起移到 GPU 或保存到权重文件。forward里直接按当前序列长度截取位置编码然后加到 embedding 上。这个实现要求d_model是偶数。如果d_model是奇数0::2和1::2的列数会不均匀赋值时会报长度不匹配。手撕时省事的方式是直接用偶数所以后面的配置都建议d_model head_dim * n_heads只要保证头数和维度能整除通常就是偶数。另一种做法是可学习位置编码self.pos_embedding nn.Parameter(torch.randn(1, max_len, d_model))两种都可以。sin/cos 的好处是不需要额外参数且对超出训练长度的序列有一定外推性可学习位置编码更灵活实现也更简单但要注意max_len限制了最大序列长度超出之后需要复制或改用其他策略。3.3 核心形状batch、seq_len、d_model 怎么对应写代码时最容易混乱的就是这三个维度。统一记成第 0 维batch_size一次处理几条样本第 1 维seq_len一条样本里的 token 数量第 2 维d_model每个 token 的向量维度后边注意力计算里还会有第 3 维和第 4 维参与例如(batch_size, n_heads, seq_len, head_dim)。如果对形状不熟每写完一个模块就打印.size()。宁可多打印几次也不要直接往下拼。4. 自注意力模块Q/K/V、缩放、多头和掩码4.1 缩放点积注意力的公式和实现自注意力可以看作一个“按相关性加权汇总”的操作。每个 token 会生成三个向量Query、Key、Value。Query 用来表达“我想找什么”Key 用来表达“我有什么”Value 是“我实际提供的信息”。先看核心函数import math import torch.nn.functional as F 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, float(-inf)) weights F.softmax(scores, dim-1) return torch.matmul(weights, v), weightsscores的 shape 是(batch_size, n_heads, seq_len_q, seq_len_k)表示每个 query 和所有 key 的相似度。除以sqrt(d_k)是为了防止点积结果过大softmax 之后梯度变得太小。如果d_k越大点积的方差越大所以这个缩放几乎是必须的。mask的 shape 要和 scores 保持一致。常见做法是传一个 4D 的 bool 张量True表示这个位置可以参与注意力False表示遮盖。被遮盖的位置在 softmax 前填成-infsoftmax 之后权重会变成 0相当于这个位置的信息完全不参与汇总。4.2 多头拆分为什么需要 view transpose多头注意力不是把模型拆成多个独立网络而是把原来的d_model维拆成n_heads个子空间每个头只看head_dim d_model // n_heads维的信息。关键代码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) self.wk nn.Linear(d_model, d_model) self.wv nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, q, k, v, maskNone): batch_size, seq_len, _ q.size() Q self.wq(q).view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2) K self.wk(k).view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2) V self.wv(v).view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2) attn, _ scaled_dot_product_attention(Q, K, V, mask) attn attn.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) return self.out_proj(attn)拆分过程分两步view(batch_size, seq_len, n_heads, head_dim)transpose(1, 2)第二行之后shape 变成(batch_size, n_heads, seq_len, head_dim)。这样每个头之间完全独立方便并行计算。最后需要把多头结果合并回(batch_size, seq_len, d_model)所以要先transpose回去再用contiguous()保证内存连续最后view合并。中间最容易报错的地方就是view之前没有contiguous()或者seq_len和d_model对不上。遇到这类问题先打印每一步 shape。4.3 mask 的两种写法mask 在注意力里有两种常见表现padding mask把无效位置遮住。比如 batch 里序列长度不一致补齐的 pad token 不应该参与注意力。因果 mask解码器里遮住未来位置让当前位置只能看自己和之前的位置。在代码实现里两者都是“把 scores 对应位置填成 -inf”。区别只在于 mask 矩阵生成方式。比如生成一个简单的因果 maskdef subsequent_mask(size): mask torch.triu(torch.ones(size, size), diagonal1).bool() return ~mask这个函数生成一个size x size的矩阵。~mask之后是下三角和主对角线为 True上三角为 False。用的时候再扩展成 4Dmask subsequent_mask(8).unsqueeze(0).unsqueeze(0)如果你在训练中遇到“生成的 token 提前看到了未来信息”十有八九是 mask 没有加或者 mask 维度不对。5. 解码器里的因果掩码和交叉注意力和编码器有什么不同5.1 encoder 的 self-attention 与 decoder 的 self-attention编码器负责理解整个源序列。它内部的 self-attention 可以看见所有位置所以不需要因果掩码。每个 token 都可以和任意其他 token 做注意力计算。解码器不一样。解码器在训练时采用 teacher forcing也就是输入一个完整的目标序列但我们不希望模型“作弊”一样提前看到后面的 token。所以在解码器第一个自注意力子层里要加因果掩码让第 i 个位置只能看到前 i 个位置。如果你写完解码器后发现训练 loss 低得离谱、但推理效果很差先检查因果掩码有没有加对。5.2 cross-attention 的 q、k、v 分别来自哪里解码器里还有一个很重要的模块cross-attention也叫 encoder-decoder attention。它的作用是让解码器在生成每一个 token 时都能从编码器输出的上下文表示里获取信息。在这个模块里Query 来自解码器上一层的输出Key 和 Value 都来自编码器输出mask 一般使用 padding mask用来屏蔽源序列里的 pad token对应代码class DecoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.cross_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) self.norm3 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) self.dropout3 nn.Dropout(dropout) def forward(self, x, encoder_output, tgt_maskNone, memory_maskNone): x self.norm1(x self.dropout1(self.self_attn(x, x, x, tgt_mask))) x self.norm2(x self.dropout2(self.cross_attn(x, encoder_output, encoder_output, memory_mask))) x self.norm3(x self.dropout3(self.ffn(x))) return x这里self_attn的 q、k、v 都是x自己所以要传三次。cross_attn的 q 是xk 和 v 都是encoder_output。如果你把 k、v 也传成x那就变成编码器输出没被用到模型很可能训不出有效结果。5.3 推理时为什么必须逐步生成训练时可以用 teacher forcing一次把整个目标序列喂进解码器。但推理时没有目标序列只能先生成第一个 token把第一个 token 接到输入序列后面再喂给解码器生成第二个 token如此反复。这就是自回归生成。代码大概是torch.no_grad() def greedy_decode(model, src, start_token0, max_len16): model.eval() src src.unsqueeze(0) memory model.encode(src) ys torch.zeros(1, 1).long() ys[0, 0] start_token for _ in range(max_len): mask subsequent_mask(ys.size(1)).unsqueeze(0).unsqueeze(0) out model.decode(memory, ys, tgt_maskmask) prob model.fc_out(out[:, -1]) next_token prob.argmax(-1).item() ys torch.cat([ys, torch.full((1, 1), next_token).long()], dim1) if next_token 2: # 假设 2 是结束符 break return ys这里有一个容易被忽略的问题每次生成新 token 时都要把整条历史序列重新输入解码器并重新计算所有之前的注意力。这不高效但逻辑简单适合初学者理解。实际生产框架会用 KV Cache 之类的优化但手撕阶段不用急着做。6. 前馈网络、残差连接、层归一化和子层排列6.1 为什么每个子层后面都要接 Add Norm注意力输出之后通常会先加一个残差再做 LayerNorm。残差让信息多一条直接通路避免深层网络里梯度消失LayerNorm 让每一层输入的量级更稳定训练更容易收敛。我在上面代码里用的是“先计算子层再加回输入最后 LayerNorm”的写法也就是 Post-Normx self.norm1(x self.dropout1(self.self_attn(x, x, x, tgt_mask)))这个顺序和原始论文一致。好处是结构简洁缺点是训练前期梯度可能不太稳定需要更小心地设置学习率。另一种更常见的工程实现是 Pre-Normx x self.dropout1(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), tgt_mask))Pre-Norm 在深层模型里更稳定训练时可以适当调大学习率。很多开源大模型都倾向于 Pre-Norm。手撕阶段建议先用 Post-Norm跑通后再改成 Pre-Norm对比一下 loss 曲线的变化。6.2 FFN 的 hidden 维度怎么选每个 Transformer 子层里还有一个逐位置的前馈网络通常叫 FeedForwardclass FeedForward(nn.Module): def __init__(self, d_model, d_ff2048, 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(F.relu(self.linear1(x))))d_ff一般取 2 到 4 倍的d_model。论文里 d_model512 时 d_ff2048。手撕时如果用 d_model32d_ff 可以取 64 或 128。d_ff太大参数会涨得很快太小非线性拟合能力会打折。注意FFN 是“逐位置”的也就是说它对序列里每个 token 施加同样的一组线性变换不跨 token 交互。这是 Transformer 里一个容易被忽略的模块但它占据了模型很大一部分参数。6.3 Post-Norm 和 Pre-Norm 的区别简单总结Post-Norm子层输出先残差最后 LayerNorm。原始论文常用训练要小心。Pre-Norm先 LayerNorm再做子层最后残差。深层模型更稳定。手撕时可先按 Post-Norm 写再改成 Pre-Norm 做对比。你会发现同样一个模型只改归一化位置loss 的收敛速度就有明显差异。这就是手撕代码的价值你能直接看到改动对训练的影响。7. 拼成完整 Transformer用“加一”玩具任务验证训练流程7.1 完整模型类前面已经写了嵌入式模块现在把它们拼成一个完整 Transformerclass Transformer(nn.Module): def __init__(self, vocab_size, d_model64, n_heads4, num_encoder_layers2, num_decoder_layers2, d_ff128, max_len5000, dropout0.1): super().__init__() self.d_model d_model self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model, max_len) self.dropout nn.Dropout(dropout) self.encoder_layers nn.ModuleList([ EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(num_encoder_layers) ]) self.decoder_layers nn.ModuleList([ DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(num_decoder_layers) ]) self.fc_out nn.Linear(d_model, vocab_size) def encode(self, src, src_maskNone): x self.dropout(self.pos_encoding(self.embedding(src) * math.sqrt(self.d_model))) for layer in self.encoder_layers: x layer(x, src_mask) return x def decode(self, memory, tgt, tgt_maskNone, memory_maskNone): x self.dropout(self.pos_encoding(self.embedding(tgt) * math.sqrt(self.d_model))) for layer in self.decoder_layers: x layer(x, memory, tgt_mask, memory_mask) return x def forward(self, src, tgt, src_maskNone, tgt_maskNone): memory self.encode(src, src_mask) out self.decode(memory, tgt, tgt_mask) return self.fc_out(out)nn.ModuleList不能换成普通 list否则模型不会把子层参数注册到模型里训练时梯度也不会更新。7.2 一个能跑通的小任务为了验证架构我用一个很简单的任务把源序列里的每个数字加一作为目标序列。比如输入[1, 2, 3]输出应该是[2, 3, 4]。这是一个有规律、但能触发完整编解码流程的任务。训练时还需要一个 EncoderLayer把我前面省略掉的编码器层补上class EncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.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) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x, maskNone): x self.norm1(x self.dropout1(self.self_attn(x, x, x, mask))) x self.norm2(x self.dropout2(self.ffn(x))) return x训练代码可以这样写import torch import torch.nn as nn model Transformer(vocab_size10, d_model64, n_heads4, num_encoder_layers2, num_decoder_layers2, d_ff128) optimizer torch.optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() for step in range(500): model.train() src torch.randint(1, 9, (4, 6)) tgt ((src) % 9) 1 tgt_input torch.zeros_like(tgt) tgt_input[:, 1:] tgt[:, :-1] tgt_output tgt mask subsequent_mask(tgt_input.size(1)).unsqueeze(0).unsqueeze(0) logits model(src, tgt_input, tgt_maskmask) loss criterion(logits.reshape(-1, 10), tgt_output.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if step % 50 0: print(step, loss.item())这个任务里vocab_size10数字范围1-90当作起始符。tgt_input是目标序列右移一位tgt_output是原始目标序列。训练时让模型根据前一个 token 预测下一个 token。7.3 训练时看什么指标只看 loss 还不够。你要看 loss 是不是稳定下降还有生成结果是不是真的符合规律。判断标准训练初期 loss 会下降但速度不一定很快。如果 loss 在某个数值附近震荡可以先降低学习率。如果 loss 一开始就卡住不降优先检查 mask、embedding 维度、tgt_input 是否右移正确。训练结束后用 greedy_decode 测试一条新序列看输出是否真的是“输入加一”。不要只盯着训练集 loss。Transformer 能力强小玩具任务很容易过拟合所以要随机生成新序列来验证泛化能力。8. 常见报错、排查顺序和调参边界8.1 维度报错的排查顺序最容易出现的错误就是形状对不上。即使你完全按论文结构写也难免在某个 view、transpose 或 cat 的地方出错。我的排查顺序是先定位报错行。PyTorch 会打印出错的具体代码行。打印相关张量 shape。比如Q.size()、scores.size()、attn.size()。从后往前检查d_model能不能被n_heads整除。检查 mask 的维度。mask 里一个常见的坑是(seq_len, seq_len)忘了扩展成 4D。检查转置之后是否调用了contiguous()。如果上次操作改变了内存布局直接view会报错。如果错误发生在matmul里最常见的两种原因是seq_len_k和seq_len_q不一致。d_k维度没有单独保留导致最后两个维度无法对齐。8.2 loss 不降和输出乱码怎么排查loss 不降不等于模型没有写对也可能是任务本身设置有问题。我一般按这个顺序排查loss 有没有变成 NaN。如果是检查学习率是否太大或者输入里有没有 NaN。去掉 mask 和不带 mask 对比一下。如果去掉 mask 之后 loss 快速下降说明因果 mask 可能遮挡错误。检查 target 里的 token 是否都在 vocab_size 范围内。nn.CrossEntropyLoss的预测维度必须覆盖 target 的每个值。检查 positional encoding 有没有加到 embedding 上。如果漏掉模型仍然能跑但效果会差。用一个小 batch 多迭代几百步确认 loss 能不能下降。如果小 batch 都不下降问题大概率在模型结构或数据构造。输出乱码通常和推理逻辑有关。最常见的问题是 greedy_decode 里没有正确使用 mask或者fc_out没有被接在 decode 后面。另外如果开始符号和结束符号定义不一致也会导致生成中断或循环。8.3 手撕 Transformer 的合理边界手撕 Transformer 适合学习、验证和理解但不适合直接上生产做大规模训练。原因很明显手写实现没有 KV Cache推理速度慢。手写实现缺少很多数值稳定性优化。手写实现没有经过大规模数据验证可能出现隐蔽 bug。如果使用成熟类库比如 Hugging Face Transformers其实已经做了很多优化和参数对齐。所以我的建议是手撕代码用于学习生产环境用成熟库。你手撕一遍之后再去读成熟库源码会顺畅很多。不要因为自己实现了完整模型就急着在任务里替换官方实现。如果想把这次练习再进一步可以尝试三件事把 Post-Norm 改成 Pre-Norm对比训练曲线。把固定 sin/cos 位置编码改成可学习位置编码看结果变化。给训练代码加上简单的早停和模型保存。这三步做完你对 Transformer 的输入、注意力、编解码器这三条主线算是真正理解了。之后再遇到模型不收敛、生成效果差、维度报错都能有一套自己的排查思路。