1. 手撕Transformer前先把输入输出的账算清楚写Transformer的PyTorch代码实现这件事我前后折腾过不下五遍。第一次照着论文把MultiHeadAttention抄完代码确实能跑loss也确实在降但你要问我输入张量里每一个维度到底代表什么、位置编码为什么非要用正弦余弦、标签为什么一定要shift一位我大概率答不上来。后来带人做项目才慢慢发现真正卡住新手的从来不是注意力公式本身——那个公式抄一遍就会了——而是输入输出这条链路上的一堆细节padding怎么处理、mask该长什么样、loss为什么算出来是nan、推理的时候为什么生成的东西前后不搭。这些问题在论文里全都被省略了因为作者默认读者是同行而我们是来学手艺的。这篇东西就是把我自己踩过的坑整理一遍围绕输入侧和输出侧这两个容易被忽略的环节展开。目标很明确给你一条从一句话到一个可训练张量的完整链路再给你一套可以复制粘贴就跑通的PyTorch代码最后附上我自己整理的排查表。适合刚学完注意力机制、准备自己动手写一版的同学也适合已经用现成nn.Transformer跑过demo、但想搞清楚内部到底发生了什么的人。读完之后你至少应该能回答三个问题一个batch的数据从文本变成张量经历了什么变换解码器的输入和标签为什么差一位以及loss里那些被mask掉的token到底去哪了。1.1 为什么建议从输入输出切进去而不是从注意力切进去注意力机制好看、好讲、好画图但它是无状态的纯计算输入输出不对它算得再对也没用。我见过太多人把注意力权重打印出来看着挺漂亮结果模型就是学不会一查是标签错位、是pad没屏蔽、是位置编码加错了维度。反过来如果你先把输入输出这条链路捋顺注意力部分你甚至可以先用nn.MultiheadAttention顶上模型照样能收敛。这就是我建议的路径先把数据的形状和数据流搞明白再回头抠注意力里的细节学习曲线会平缓很多。另外一个原因是Transformer原始论文是编码器-解码器结构但今天大家用的很多是只有编码器的BERT类或者只有解码器的GPT类。这两种变体在输入输出上的差别非常大编码器只需要一个padding mask解码器需要因果mask加cross-attention的memory mask训练和推理时输入还不一样。如果你直接上来啃多头注意力的矩阵乘法很容易把这些差异混在一起最后写出来的代码训练能跑、推理直接崩。所以先把账算清楚是省时间而不是绕远路。1.2 一份能跑通的最小闭环长什么样我习惯把整条链路拆成六段来记文本切分得到token id序列、padding对齐成定长、Embedding查表得到稠密向量、加上位置编码注入顺序信息、送进编码器或解码器做特征变换、最后接一个线性层投影回词表大小做分类。这条链路里第二段和第六段是最容易出问题的前者涉及pad的处理后者涉及标签和维度的对齐。闭环能不能跑通我一般用一个特别土的办法验证拿一个序列反转任务当试金石。输入一串随机token让模型输出它的逆序。这个任务有几个好处——词表小、序列短、规律明确模型如果结构正确几百步就该学到接近百分之百的准确率如果结构有错它连这个都学不会。比起直接上机器翻译数据集用反转任务定位问题快得多。下面所有代码都是围绕这个任务来验证的你把它换成真实数据只需要改数据加载部分。2. 输入侧从文本到张量的整条链路输入侧的核心矛盾是变长和定长的矛盾。自然语言长度不固定但GPU上的张量必须是规整的矩形。解决方式是padding也就是拿一个特殊符号把短句子补齐让整个batch的序列长度一致。这个操作本身很简单但它带来的连锁反应贯穿整个模型Embedding层要忽略pad、注意力要屏蔽pad、loss要忽略pad。三处有任何一处漏了模型都会学歪。2.1 词表和token id先把数字这件事说清楚词表是输入侧的起点。它的本质是一个字符串到整数的映射表加上两个方向的操作encode把文本变成id列表decode把id列表变回文本。真实项目里这一步通常交给现成的分词器但手撕的时候我建议你自己维护一个极简版因为你需要对特殊符号的位置有完全的掌控。下面这个类我用了很久够简单也够用。要注意的是四个特殊符号的约定PAD必须是0因为后面np/pytorch里很多默认参数都跟0对齐UNK用来兜住词表外的词BOS和EOS标记序列的起止其中BOS只在解码器输入里出现EOS只在标签里出现——这个不对称是后面shift right的根源。class Vocab: PAD, UNK, BOS, EOS 0, 1, 2, 3 def __init__(self, tokens): # 特殊符号先占位保证 id 固定 self.itos [pad, unk, bos, eos] self.stoi {t: i for i, t in enumerate(self.itos)} for t in tokens: if t not in self.stoi: self.stoi[t] len(self.itos) self.itos.append(t) def __len__(self): return len(self.itos) def encode(self, tokens, add_bosFalse, add_eosFalse): ids [self.stoi.get(t, self.UNK) for t in tokens] if add_bos: ids [self.BOS] ids if add_eos: ids ids [self.EOS] return ids有个细节值得单独说特殊符号的id一定要在构建词表之前就固定下来不能等到最后再append。我第一版代码就是先把语料里的词加进去、最后才补特殊符号结果每换一份语料特殊符号的id都在变之前存的检查点全部作废。这种错误不会报错只会在推理的时候给你输出一堆莫名其妙的结果排查起来极其难受。2.2 Embedding层的三个参数陷阱Embedding本质就是一张查找表形状是(词表大小, d_model)。看上去没什么可讲的但有三个地方特别容易出问题。第一个是padding_idx。这个参数告诉PyTorch哪一行是填充位它的效果是前向传播时该位置的输出恒为0反向传播时该行的梯度永远是0。很多人以为设了这个参数就能自动屏蔽pad其实不是——它只影响Embedding这一层自己的输出和梯度注意力里的pad位置照样会参与计算照样会通过softmax影响其他位置。所以padding_idx只是第一道防线mask才是主防线。第二个是缩放。原始论文里Embedding的输出要乘以sqrt(d_model)原因是位置编码的数值范围大致在[-1,1]而Embedding初始化后的方差约为1/d_model量级差了一个数量级。如果不缩放位置编码会淹没词向量信息模型学到的表示会偏向位置而忽略语义。这个乘法看起来微不足道但我在小模型上做过对比不加缩放的版本收敛明显更慢。代码上就一行x self.embed(x) * math.sqrt(self.d_model)。第三个是初始化。默认初始化在小词表上问题不大但词表上万之后建议用正态分布手动初始化并把padding_idx那行清零。理由是正态分布能让不同词向量的初始余弦相似度接近0避免一开始就有大量词挤在一起训练初期梯度方向更干净。这个技巧在大词表场景下收益比较明显。class TokenEmbedding(nn.Module): def __init__(self, vocab_size, d_model, pad_idx0): super().__init__() self.d_model d_model self.embed nn.Embedding(vocab_size, d_model, padding_idxpad_idx) nn.init.normal_(self.embed.weight, mean0.0, stdd_model ** -0.5) with torch.no_grad(): self.embed.weight[pad_idx].fill_(0) def forward(self, x): # x: (B, L) - (B, L, d_model) return self.embed(x) * math.sqrt(self.d_model)注意padding_idx的梯度清零是永久的如果你后面把pad位置替换成别的用途比如用它做padding mask的载体记得那一行权重永远是0不要再指望它承载信息。2.3 位置编码公式简单实现全是坑位置编码是输入侧最数学的部分。它要解决的问题是注意力本身是置换不变的打乱输入顺序输出只是跟着打乱模型无法感知谁在前谁在后。所以必须人为地把位置信息注入进去。原始方案是用不同频率的正弦余弦函数生成一组固定的向量直接加到词向量上。为什么用三角函数而不是直接用一个可学习的位置Embedding我自己的理解有两层一是三角函数具有相对位置可表示的性质位置pk的编码可以表示成位置p编码的线性变换模型更容易学到相对距离这个概念二是外推性好训练时见过的最大长度是512推理时来一个800长度的序列三角函数照样能算出合理的编码而可学习的位置Embedding在那个位置上是没有训练过的。但也有代价三角函数是固定的不参与训练所以它未必是该任务下最优的位置表示。后来很多工作改成了相对位置编码或者旋转位置编码效果确实更好。不过作为手撕的第一版三角函数是最合适的选择——它不引入额外参数出错的地方少便于定位问题。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() assert d_model % 2 0, d_model 建议取偶数否则 sin/cos 交错赋值会长度不匹配 self.dropout nn.Dropout(dropout) pe torch.zeros(max_len, d_model) pos torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # (max_len, 1) # 1 / 10000^(2i/d_model)在对数空间计算防止溢出 div torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(pos * div) # 偶数维 pe[:, 1::2] torch.cos(pos * div) # 奇数维 # (1, max_len, d_model)注册成 buffer随模型一起搬到 GPU但不参与训练 self.register_buffer(pe, pe.unsqueeze(0)) def forward(self, x): # x: (B, L, d_model) x x self.pe[:, :x.size(1)] return self.dropout(x)这段代码里有三个我踩过的坑。第一个是register_buffer和nn.Parameter的区别位置编码不训练所以不能用Parameter但你必须用buffer注册否则.to(device)搬不走它训练时就会出现CPU张量和GPU张量相加的报错。第二个是d_model取奇数时pe[:, 0::2]和pe[:, 1::2]的长度会差一个形状对不上直接抛异常所以我加了断言。第三个是dropout的位置论文里是在加完位置编码之后做dropout这个顺序不能反先dropout再加法等于没对位置做正则。顺带说一句max_len的设定。这个值决定了预计算矩阵的大小设太大浪费显存5000乘512也就10M其实还好设太小推理时序列一长就索引越界。我的习惯是按业务最长序列的两倍来设留足余量。如果你确实需要超长序列与其把max_len调到十万不如考虑换成旋转位置编码那套方案它是按需计算的不依赖预计算表。3. 输出侧标签、mask与loss的三方对账输出侧是新手翻车重灾区。常见症状包括loss从第一步开始就不降、loss降到某个值就卡住不动、训练集准确率很高但生成结果全是重复词。这些问题八成不是模型结构的问题而是标签、mask、loss这三者之间没有对齐。3.1 shift right解码器输入和标签为什么差一位语言模型训练时用的是教师强制策略也就是说不管模型上一步预测得对不对下一步喂给它的都是真实的上一个token。这样做的目的是让训练稳定、可以并行——因为所有位置的输入都已知不需要等前一个位置算完。具体来说如果目标序列是[BOS, w1, w2, w3, EOS]那么解码器的输入是[BOS, w1, w2, w3]对应的标签是[w1, w2, w3, EOS]。两者长度相同但整体错开一位。这个错位在代码里就是两行切片# tgt: (B, Lt)形如 [BOS, w1, w2, w3, EOS] tgt_in tgt[:, :-1] # 解码器输入: [BOS, w1, w2, w3] tgt_out tgt[:, 1:] # 监督标签: [w1, w2, w3, EOS]为什么要这么设计因为模型在第t个位置上的输出预测的应该是给定前t个token之后的下一个token也就是标签序列的第t个元素。这样每个位置的损失函数就是标准的多分类交叉熵可以直接用GPU并行计算。我第一次写的时候没做这个错位输入和标签完全相同。结果模型很快学会了一个偷懒的解——把输入原样复制到输出loss确实降了但推理时因为看不到下一个token生成的序列会完全乱掉。这种错误在训练曲线上看不出来只有推理时才暴露非常阴险。3.2 三种mask的生成逻辑与形状mask是输出侧的核心。它的作用是在softmax之前把不该关注的位置的分数设成负无穷这样softmax之后那些位置的权重就趋近于0。用加法还是用masked_fill都行我习惯用布尔类型的masked_fill语义更清楚。Transformer里有三种mask很多人把它们混在一起我列个表对照一下。mask名称作用位置形状生成规则屏蔽目的src_mask编码器自注意力(B, 1, Ls) 或 (B, Ls, Ls)src ! PAD不让pad参与注意力tgt_mask解码器自注意力(B, Lt, Lt)因果下三角 且 tgt ! PAD防未来信息泄露 屏蔽padmemory_mask解码器交叉注意力(B, 1, Ls)src ! PAD不让解码器关注到pad这三者的生成方式如下。注意布尔mask里True表示保留False表示屏蔽和有些教程用1/0的约定相反你要统一好别混着用。def make_pad_mask(seq, pad_idx0): (B, L) - (B, 1, L)True 表示是真实token return (seq ! pad_idx).unsqueeze(1) def make_causal_mask(L, device): (L, L) 下三角矩阵True 表示可见 return torch.tril(torch.ones(L, L, dtypetorch.bool, devicedevice)) # 组合先算 pad mask再和因果 mask 做广播与运算 src_mask make_pad_mask(src) # (B, 1, Ls) tgt_pad make_pad_mask(tgt_in) # (B, 1, Lt) causal make_causal_mask(tgt_in.size(1), tgt_in.device).unsqueeze(0) # (1, Lt, Lt) tgt_mask tgt_pad causal # (B, Lt, Lt)这里有个形状上的坑tgt_pad是(B, 1, Lt)causal是(1, Lt, Lt)两者做之后广播成(B, Lt, Lt)。广播的结果正好是注意力分数的形状可以直接用。关于因果mask还有个细节它是在每个batch里都是同一份所以完全可以缓存起来不必每次forward都重新生成。序列长度固定的话这个缓存能省下一点开销在长序列上比较可观。提示屏蔽用的是负无穷而不是0。如果你用0去填充被屏蔽的位置softmax之后那些位置仍然会分到正权重只是权重小一点而已pad信息照样会漏进来。这个错误不会报错只会让模型悄悄地学会利用pad。3.3 loss计算ignore_index和log_softmax的关系输出侧的最后一环是损失函数。模型的最后一层是一个线性层把d_model维投影回词表大小得到每个位置的logits。注意这里是logits不是概率不要手贱加softmax——因为CrossEntropyLoss内部已经包含了log_softmax你加了就重复了梯度会异常。class Generator(nn.Module): def __init__(self, d_model, vocab_size): super().__init__() self.proj nn.Linear(d_model, vocab_size) def forward(self, x): # x: (B, L, d_model) - (B, L, V) return self.proj(x)loss计算时需要把(B, L, V)拉平成(B*L, V)标签拉平成(B*L)这是CrossEntropyLoss期望的形状也就是NLLLoss的形状约定。关键是ignore_index参数要设成PAD的id这样被padding的位置不产生损失也不会贡献梯度。criterion nn.CrossEntropyLoss(ignore_index0, label_smoothing0.1) logits model(src, tgt_in, src_mask, tgt_mask) # (B, Lt, V) loss criterion( logits.reshape(-1, logits.size(-1)), # (B*Lt, V) tgt_out.reshape(-1) # (B*Lt,) )label_smoothing这个参数我建议加上尤其是小数据集上。它的作用是把硬标签从1和0软化成0.9和0.1/(V-1)左右防止模型对某一个类别过度自信泛化性会好一些。代价是训练loss数值会变大因为熵本身就高了所以不要拿加了平滑的loss和没加的对比会误判。还有一个常见疑问为什么标签要用tgt_out而不是tgt_in回到3.1那节因为第t个位置的输出预测的是下一个token所以标签必须右移一位。你可以在代码里打印前几条样本确认一下肉眼对一遍比反复读论文快。4. 完整实现用反转任务把闭环跑通前面拆的都是零件这一节把它们拼起来。我用序列反转当验证任务投入最小、反馈最快。整个模型包括三个部分编码器层、解码器层、以及前面讲的输入输出模块。4.1 多头注意力与编码解码层先写多头注意力。这里的核心是把(B, L, d_model)拆成(B, h, L, d_k)。用view加transpose是最省事的写法但要注意view之前张量必须是连续的transpose之后就不连续了所以反向reshape回原形状时要先contiguous()。class MHA(nn.Module): def __init__(self, d_model, n_head, dropout0.1): super().__init__() assert d_model % n_head 0 self.h n_head self.d_k d_model // n_head 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.w_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, q, k, v, maskNone): B, Lq q.size(0), q.size(1) Lk k.size(1) q self.w_q(q).view(B, Lq, self.h, self.d_k).transpose(1, 2) k self.w_k(k).view(B, Lk, self.h, self.d_k).transpose(1, 2) v self.w_v(v).view(B, Lk, self.h, self.d_k).transpose(1, 2) scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(~mask, float(-inf)) attn torch.softmax(scores, dim-1) attn self.dropout(attn) out torch.matmul(attn, v) # (B, h, Lq, d_k) out out.transpose(1, 2).contiguous().view(B, Lq, -1) return self.w_o(out)编码器和解码器层就是注意力加前馈网络每个子层外面套一层残差和LayerNorm。这里我用的是最传统的Post-LN先残差后Norm对新手更直观。如果你要训练深层模型建议换成Pre-LN因为Pre-LN不需要在残差分支上加warmup训练稳定性更好。class FFN(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.net nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) def forward(self, x): return self.net(x) class EncoderLayer(nn.Module): def __init__(self, d_model, n_head, d_ff, dropout0.1): super().__init__() self.attn MHA(d_model, n_head, dropout) self.ffn FFN(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.drop nn.Dropout(dropout) def forward(self, x, mask): x self.norm1(x self.drop(self.attn(x, x, x, mask))) x self.norm2(x self.drop(self.ffn(x))) return x class DecoderLayer(nn.Module): def __init__(self, d_model, n_head, d_ff, dropout0.1): super().__init__() self.self_attn MHA(d_model, n_head, dropout) self.cross_attn MHA(d_model, n_head, dropout) self.ffn FFN(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.drop nn.Dropout(dropout) def forward(self, x, memory, tgt_mask, mem_mask): x self.norm1(x self.drop(self.self_attn(x, x, x, tgt_mask))) x self.norm2(x self.drop(self.cross_attn(x, memory, memory, mem_mask))) x self.norm3(x self.drop(self.ffn(x))) return x交叉注意力那一层的mask要注意它屏蔽的是key和value所在的源序列长度所以形状必须是(B, 1, Ls)而不是目标长度。我第一次写的时候顺手把tgt_mask传进去了seq长度一样时不会报错长度不一样时直接形状不匹配或者更糟——广播成功但屏蔽了错误的位置。4.2 数据构造与训练循环反转任务的构造很简单随机生成一串token目标是它的逆序。这里我用torch.randint直接生成不需要任何外部数据文件方便你复制过去就运行。PAD, BOS, EOS, VOCAB_SIZE 0, 1, 2, 50 def make_batch(batch_size32, seq_len8): # 生成范围 [3, VOCAB_SIZE) 的token避开特殊符号 src torch.randint(3, VOCAB_SIZE, (batch_size, seq_len)) rev torch.flip(src, dims[1]) bos torch.full((batch_size, 1), BOS) eos torch.full((batch_size, 1), EOS) tgt torch.cat([bos, rev, eos], dim1) # [BOS, rev..., EOS] return src, tgt训练循环里我用了Noam式学习率也就是线性warmup加逆平方根衰减。它对Transformer这种深层结构几乎是标配因为训练初期梯度方向很不稳定学习率太大会直接发散。这个调度器不需要手调warmup_steps设成总步数的十分之一左右就够了。def train(steps800, d_model64, n_head4, n_layer2, d_ff128): model Seq2Seq(d_model, n_head, n_layer, d_ff, VOCAB_SIZE).to(DEVICE) opt torch.optim.Adam(model.parameters(), lr0.0, betas(0.9, 0.98), eps1e-9) criterion nn.CrossEntropyLoss(ignore_indexPAD, label_smoothing0.0) for step in range(1, steps 1): # warmup 逆平方根衰减 lr d_model ** -0.5 * min(step ** -0.5, step * 400 ** -1.5) for g in opt.param_groups: g[lr] lr src, tgt make_batch() src, tgt src.to(DEVICE), tgt.to(DEVICE) tgt_in, tgt_out tgt[:, :-1], tgt[:, 1:] src_mask make_pad_mask(src) tgt_mask make_pad_mask(tgt_in) make_causal_mask(tgt_in.size(1), DEVICE).unsqueeze(0) logits model(src, tgt_in, src_mask, tgt_mask) loss criterion(logits.reshape(-1, logits.size(-1)), tgt_out.reshape(-1)) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() if step % 100 0: print(fstep {step:4d} | lr {lr:.2e} | loss {loss.item():.4f})梯度裁剪那行不要省。Transformer的梯度在训练初期会非常激进尤其是最后的投影层偶尔会冒出一个很大的梯度把参数带偏。裁剪阈值1.0是经验值一般不用改。如果发现loss偶尔出现尖刺然后又恢复八成就是没裁剪。跑起来之后你应该看到这样的曲线前一百步loss从4左右降到0.5以下三百步左右跌到接近0因为词表只有50、序列长度只有8任务非常容易。如果到了五百步loss还在1.5以上那基本可以确定是mask或者标签那里出了错不要怀疑模型容量去检查shift。4.3 推理阶段自回归生成训练时的输入输出都是已知的推理时不是。推理要从一个BOS开始一步步生成每生成一个token就把它拼接到输入末尾再让模型预测下一个。这个过程叫自回归解码。torch.no_grad() def greedy_decode(model, src, max_len20): model.eval() src_mask make_pad_mask(src) memory model.encode(src, src_mask) B src.size(0) ys torch.full((B, 1), BOS, dtypetorch.long, devicesrc.device) finished torch.zeros(B, dtypetorch.bool, devicesrc.device) for _ in range(max_len): tgt_mask make_pad_mask(ys) make_causal_mask(ys.size(1), ys.device).unsqueeze(0) out model.decode(memory, ys, tgt_mask, src_mask) logits model.generator(out[:, -1]) # 只取最后一个位置 next_id logits.argmax(dim-1, keepdimTrue) # (B, 1) ys torch.cat([ys, next_id], dim1) finished | (next_id.squeeze(1) EOS) if finished.all(): break return ys这段代码有几个点值得强调。第一model.eval()和torch.no_grad()必须加前者关掉dropout后者省显存。第二每次循环都要重新计算tgt_mask因为序列长度在变因果mask的尺寸也跟着变。第三只需要取最后一个位置的logits因为前面的位置早就在上一步生成过了重复计算是浪费——这也是KV缓存要解决的问题。第四EOS判断是逐样本的一个batch里有的样本先结束理论上可以给它补PAD并跳过计算但为了代码简单我先没做。如果你要上生产这里至少有两个优化空间。一是KV缓存把每层注意力的key和value缓存下来每一步只算新token复杂度从O(L²)降到O(L)。二是批量采样策略贪心解码容易生成重复内容换成beam search或者top-k采样会自然很多。这两块展开又是一篇这里先提一句知道方向就行。5. 排查实录我踩过的坑和对应解法这部分是全篇最实用的。前面讲的是应该怎么做这里讲的是做错了会怎样、怎么找回来。我把遇到的问题按症状分类整理你可以当速查表用。5.1 形状与类型类问题速查表形状错误是最好查的因为PyTorch会直接报错。麻烦的是那些不报错但结果悄悄错的我单独在后面列。报错/症状大概率原因定位方法matmul维度不匹配mask形状与scores不匹配在masked_fill前打印mask和scores的shapemasked_fill dtype错误mask用了float不是bool统一用(seq ! PAD)生成得到bool位置编码相加报错pe没注册成buffer还在CPU检查register_buffer是否正确调用传给交叉注意力的mask长度不对用了tgt_mask而非src_mask交叉注意力的mask长度必须等于源序列长度CrossEntropyLoss维度错误logits没reshape成二维用.reshape(-1, V)压平推理时越界max_len超过了位置编码表长度检查max_len设置我这个表格里最常出现的是第二条。因为很多老教程用的是加性mask0和负无穷生成的mask是float类型masked_fill要求布尔两者混用就会报错。我的建议是全程只用布尔mask加性mask如果需要比如实现相对位置偏置再单独处理。5.2 训练不收敛的几种典型表现第一种loss从第一步就在4.0附近不动或者直接是nan。这种情况先查学习率warmup没做好或者初始lr太大是最常见的。其次查标签里有没有出现超出词表范围的id——CrossEntropyLoss在标签越界时会直接给nan而不是报错非常坑。我养成的习惯是每次造完数据加一句断言assert tgt.max() vocab_size。第二种loss能降到2左右但卡住不动。这个位置恰好对应均匀分布的熵说明模型学到的就是词频分布没有任何结构信息。原因通常是mask写反了把该保留的位置屏蔽了、该屏蔽的位置保留了。特别是因果mask如果上下三角搞反模型只能看到未来的token当然学不出东西。验证方法很简单把mask矩阵打印出来看下三角是不是True。第三种训练loss很低但验证集或者推理结果很差。这就是典型的过拟合或训练推理不一致。如果不一致去检查三点推理时有没有加因果mask、有没有误用dropout、EOS有没有正确终止。我遇到过一次是推理时忘了把模型切到eval模式dropout还在生效输出全是乱的。第四种loss正常但在某个值附近震荡不下降。这时候把学习率调小十倍试试如果还是不行检查一下label_smoothing——加了平滑之后loss有下限别把它当成不收敛。提示每次调参只改一个变量。我早期吃过这个亏一次改三处效果变好了也不知道是哪一处起的作用下次遇到问题还是不知道怎么复现。5.3 几个我希望早点知道的经验关于padding_idx前面说过它只保护Embedding自己。但还有个隐藏行为如果你在位置编码里对pad位置也加了位置向量那pad位置就不再是全零了它会变成一个纯位置向量依然会污染注意力。解决办法很简单做完位置编码后把pad位置乘回一个掩码x x.masked_fill((tgt_in PAD).unsqueeze(-1), 0.0)关于batch_firstPyTorch原生的nn.Transformer默认是batch_firstFalse也就是序列维在前。这是历史遗留问题很多人包括我第一次用的时候根本没注意模型跑起来了但结果全错。现在的版本支持传batch_firstTrue我强烈建议打开它整条链路都用(B, L, d)的约定不用来回转置。手写的话就更简单了一开始就坚持(B, L, d)别给自己添麻烦。关于显存注意力矩阵是O(L²)的序列长度从512涨到1024显存直接翻四倍。手撕版本在调试阶段很容易OOM我的做法是先把batch和序列长度压到很小比如批量8、长度16把逻辑跑通再加规模。别一上来就上真实数据集的规模调不完的。关于可复现性训练前一定固定随机种子包括torch.manual_seed、random.seed、numpy.random.seed并且注意CuDNN的确定性选项。我吃过一次亏同一份代码跑三遍得到三个不同的收敛速度花了半天时间怀疑人生最后发现是随机种子没固定。这个习惯在小规模调试阶段尤其重要不然你根本分不清是改代码起了作用还是运气好。最后说个心态上的事。手撕Transformer的意义不在于以后真的手写而在于你能读懂报错的含义、能判断一个开源实现哪里不对劲、能在模型不工作时快速定位到是数据、mask还是loss的问题。我见过太多人调包调得很顺一遇到形状不匹配就束手无策。多写两遍把这条链路刻进肌肉记忆里后面看任何变体——无论它是改注意力、改位置编码还是改归一化位置——你心里都会有一个稳定的坐标系。
手撕Transformer：PyTorch输入输出、Mask与Loss全链路
1. 手撕Transformer前先把输入输出的账算清楚写Transformer的PyTorch代码实现这件事我前后折腾过不下五遍。第一次照着论文把MultiHeadAttention抄完代码确实能跑loss也确实在降但你要问我输入张量里每一个维度到底代表什么、位置编码为什么非要用正弦余弦、标签为什么一定要shift一位我大概率答不上来。后来带人做项目才慢慢发现真正卡住新手的从来不是注意力公式本身——那个公式抄一遍就会了——而是输入输出这条链路上的一堆细节padding怎么处理、mask该长什么样、loss为什么算出来是nan、推理的时候为什么生成的东西前后不搭。这些问题在论文里全都被省略了因为作者默认读者是同行而我们是来学手艺的。这篇东西就是把我自己踩过的坑整理一遍围绕输入侧和输出侧这两个容易被忽略的环节展开。目标很明确给你一条从一句话到一个可训练张量的完整链路再给你一套可以复制粘贴就跑通的PyTorch代码最后附上我自己整理的排查表。适合刚学完注意力机制、准备自己动手写一版的同学也适合已经用现成nn.Transformer跑过demo、但想搞清楚内部到底发生了什么的人。读完之后你至少应该能回答三个问题一个batch的数据从文本变成张量经历了什么变换解码器的输入和标签为什么差一位以及loss里那些被mask掉的token到底去哪了。1.1 为什么建议从输入输出切进去而不是从注意力切进去注意力机制好看、好讲、好画图但它是无状态的纯计算输入输出不对它算得再对也没用。我见过太多人把注意力权重打印出来看着挺漂亮结果模型就是学不会一查是标签错位、是pad没屏蔽、是位置编码加错了维度。反过来如果你先把输入输出这条链路捋顺注意力部分你甚至可以先用nn.MultiheadAttention顶上模型照样能收敛。这就是我建议的路径先把数据的形状和数据流搞明白再回头抠注意力里的细节学习曲线会平缓很多。另外一个原因是Transformer原始论文是编码器-解码器结构但今天大家用的很多是只有编码器的BERT类或者只有解码器的GPT类。这两种变体在输入输出上的差别非常大编码器只需要一个padding mask解码器需要因果mask加cross-attention的memory mask训练和推理时输入还不一样。如果你直接上来啃多头注意力的矩阵乘法很容易把这些差异混在一起最后写出来的代码训练能跑、推理直接崩。所以先把账算清楚是省时间而不是绕远路。1.2 一份能跑通的最小闭环长什么样我习惯把整条链路拆成六段来记文本切分得到token id序列、padding对齐成定长、Embedding查表得到稠密向量、加上位置编码注入顺序信息、送进编码器或解码器做特征变换、最后接一个线性层投影回词表大小做分类。这条链路里第二段和第六段是最容易出问题的前者涉及pad的处理后者涉及标签和维度的对齐。闭环能不能跑通我一般用一个特别土的办法验证拿一个序列反转任务当试金石。输入一串随机token让模型输出它的逆序。这个任务有几个好处——词表小、序列短、规律明确模型如果结构正确几百步就该学到接近百分之百的准确率如果结构有错它连这个都学不会。比起直接上机器翻译数据集用反转任务定位问题快得多。下面所有代码都是围绕这个任务来验证的你把它换成真实数据只需要改数据加载部分。2. 输入侧从文本到张量的整条链路输入侧的核心矛盾是变长和定长的矛盾。自然语言长度不固定但GPU上的张量必须是规整的矩形。解决方式是padding也就是拿一个特殊符号把短句子补齐让整个batch的序列长度一致。这个操作本身很简单但它带来的连锁反应贯穿整个模型Embedding层要忽略pad、注意力要屏蔽pad、loss要忽略pad。三处有任何一处漏了模型都会学歪。2.1 词表和token id先把数字这件事说清楚词表是输入侧的起点。它的本质是一个字符串到整数的映射表加上两个方向的操作encode把文本变成id列表decode把id列表变回文本。真实项目里这一步通常交给现成的分词器但手撕的时候我建议你自己维护一个极简版因为你需要对特殊符号的位置有完全的掌控。下面这个类我用了很久够简单也够用。要注意的是四个特殊符号的约定PAD必须是0因为后面np/pytorch里很多默认参数都跟0对齐UNK用来兜住词表外的词BOS和EOS标记序列的起止其中BOS只在解码器输入里出现EOS只在标签里出现——这个不对称是后面shift right的根源。class Vocab: PAD, UNK, BOS, EOS 0, 1, 2, 3 def __init__(self, tokens): # 特殊符号先占位保证 id 固定 self.itos [pad, unk, bos, eos] self.stoi {t: i for i, t in enumerate(self.itos)} for t in tokens: if t not in self.stoi: self.stoi[t] len(self.itos) self.itos.append(t) def __len__(self): return len(self.itos) def encode(self, tokens, add_bosFalse, add_eosFalse): ids [self.stoi.get(t, self.UNK) for t in tokens] if add_bos: ids [self.BOS] ids if add_eos: ids ids [self.EOS] return ids有个细节值得单独说特殊符号的id一定要在构建词表之前就固定下来不能等到最后再append。我第一版代码就是先把语料里的词加进去、最后才补特殊符号结果每换一份语料特殊符号的id都在变之前存的检查点全部作废。这种错误不会报错只会在推理的时候给你输出一堆莫名其妙的结果排查起来极其难受。2.2 Embedding层的三个参数陷阱Embedding本质就是一张查找表形状是(词表大小, d_model)。看上去没什么可讲的但有三个地方特别容易出问题。第一个是padding_idx。这个参数告诉PyTorch哪一行是填充位它的效果是前向传播时该位置的输出恒为0反向传播时该行的梯度永远是0。很多人以为设了这个参数就能自动屏蔽pad其实不是——它只影响Embedding这一层自己的输出和梯度注意力里的pad位置照样会参与计算照样会通过softmax影响其他位置。所以padding_idx只是第一道防线mask才是主防线。第二个是缩放。原始论文里Embedding的输出要乘以sqrt(d_model)原因是位置编码的数值范围大致在[-1,1]而Embedding初始化后的方差约为1/d_model量级差了一个数量级。如果不缩放位置编码会淹没词向量信息模型学到的表示会偏向位置而忽略语义。这个乘法看起来微不足道但我在小模型上做过对比不加缩放的版本收敛明显更慢。代码上就一行x self.embed(x) * math.sqrt(self.d_model)。第三个是初始化。默认初始化在小词表上问题不大但词表上万之后建议用正态分布手动初始化并把padding_idx那行清零。理由是正态分布能让不同词向量的初始余弦相似度接近0避免一开始就有大量词挤在一起训练初期梯度方向更干净。这个技巧在大词表场景下收益比较明显。class TokenEmbedding(nn.Module): def __init__(self, vocab_size, d_model, pad_idx0): super().__init__() self.d_model d_model self.embed nn.Embedding(vocab_size, d_model, padding_idxpad_idx) nn.init.normal_(self.embed.weight, mean0.0, stdd_model ** -0.5) with torch.no_grad(): self.embed.weight[pad_idx].fill_(0) def forward(self, x): # x: (B, L) - (B, L, d_model) return self.embed(x) * math.sqrt(self.d_model)注意padding_idx的梯度清零是永久的如果你后面把pad位置替换成别的用途比如用它做padding mask的载体记得那一行权重永远是0不要再指望它承载信息。2.3 位置编码公式简单实现全是坑位置编码是输入侧最数学的部分。它要解决的问题是注意力本身是置换不变的打乱输入顺序输出只是跟着打乱模型无法感知谁在前谁在后。所以必须人为地把位置信息注入进去。原始方案是用不同频率的正弦余弦函数生成一组固定的向量直接加到词向量上。为什么用三角函数而不是直接用一个可学习的位置Embedding我自己的理解有两层一是三角函数具有相对位置可表示的性质位置pk的编码可以表示成位置p编码的线性变换模型更容易学到相对距离这个概念二是外推性好训练时见过的最大长度是512推理时来一个800长度的序列三角函数照样能算出合理的编码而可学习的位置Embedding在那个位置上是没有训练过的。但也有代价三角函数是固定的不参与训练所以它未必是该任务下最优的位置表示。后来很多工作改成了相对位置编码或者旋转位置编码效果确实更好。不过作为手撕的第一版三角函数是最合适的选择——它不引入额外参数出错的地方少便于定位问题。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() assert d_model % 2 0, d_model 建议取偶数否则 sin/cos 交错赋值会长度不匹配 self.dropout nn.Dropout(dropout) pe torch.zeros(max_len, d_model) pos torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # (max_len, 1) # 1 / 10000^(2i/d_model)在对数空间计算防止溢出 div torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(pos * div) # 偶数维 pe[:, 1::2] torch.cos(pos * div) # 奇数维 # (1, max_len, d_model)注册成 buffer随模型一起搬到 GPU但不参与训练 self.register_buffer(pe, pe.unsqueeze(0)) def forward(self, x): # x: (B, L, d_model) x x self.pe[:, :x.size(1)] return self.dropout(x)这段代码里有三个我踩过的坑。第一个是register_buffer和nn.Parameter的区别位置编码不训练所以不能用Parameter但你必须用buffer注册否则.to(device)搬不走它训练时就会出现CPU张量和GPU张量相加的报错。第二个是d_model取奇数时pe[:, 0::2]和pe[:, 1::2]的长度会差一个形状对不上直接抛异常所以我加了断言。第三个是dropout的位置论文里是在加完位置编码之后做dropout这个顺序不能反先dropout再加法等于没对位置做正则。顺带说一句max_len的设定。这个值决定了预计算矩阵的大小设太大浪费显存5000乘512也就10M其实还好设太小推理时序列一长就索引越界。我的习惯是按业务最长序列的两倍来设留足余量。如果你确实需要超长序列与其把max_len调到十万不如考虑换成旋转位置编码那套方案它是按需计算的不依赖预计算表。3. 输出侧标签、mask与loss的三方对账输出侧是新手翻车重灾区。常见症状包括loss从第一步开始就不降、loss降到某个值就卡住不动、训练集准确率很高但生成结果全是重复词。这些问题八成不是模型结构的问题而是标签、mask、loss这三者之间没有对齐。3.1 shift right解码器输入和标签为什么差一位语言模型训练时用的是教师强制策略也就是说不管模型上一步预测得对不对下一步喂给它的都是真实的上一个token。这样做的目的是让训练稳定、可以并行——因为所有位置的输入都已知不需要等前一个位置算完。具体来说如果目标序列是[BOS, w1, w2, w3, EOS]那么解码器的输入是[BOS, w1, w2, w3]对应的标签是[w1, w2, w3, EOS]。两者长度相同但整体错开一位。这个错位在代码里就是两行切片# tgt: (B, Lt)形如 [BOS, w1, w2, w3, EOS] tgt_in tgt[:, :-1] # 解码器输入: [BOS, w1, w2, w3] tgt_out tgt[:, 1:] # 监督标签: [w1, w2, w3, EOS]为什么要这么设计因为模型在第t个位置上的输出预测的应该是给定前t个token之后的下一个token也就是标签序列的第t个元素。这样每个位置的损失函数就是标准的多分类交叉熵可以直接用GPU并行计算。我第一次写的时候没做这个错位输入和标签完全相同。结果模型很快学会了一个偷懒的解——把输入原样复制到输出loss确实降了但推理时因为看不到下一个token生成的序列会完全乱掉。这种错误在训练曲线上看不出来只有推理时才暴露非常阴险。3.2 三种mask的生成逻辑与形状mask是输出侧的核心。它的作用是在softmax之前把不该关注的位置的分数设成负无穷这样softmax之后那些位置的权重就趋近于0。用加法还是用masked_fill都行我习惯用布尔类型的masked_fill语义更清楚。Transformer里有三种mask很多人把它们混在一起我列个表对照一下。mask名称作用位置形状生成规则屏蔽目的src_mask编码器自注意力(B, 1, Ls) 或 (B, Ls, Ls)src ! PAD不让pad参与注意力tgt_mask解码器自注意力(B, Lt, Lt)因果下三角 且 tgt ! PAD防未来信息泄露 屏蔽padmemory_mask解码器交叉注意力(B, 1, Ls)src ! PAD不让解码器关注到pad这三者的生成方式如下。注意布尔mask里True表示保留False表示屏蔽和有些教程用1/0的约定相反你要统一好别混着用。def make_pad_mask(seq, pad_idx0): (B, L) - (B, 1, L)True 表示是真实token return (seq ! pad_idx).unsqueeze(1) def make_causal_mask(L, device): (L, L) 下三角矩阵True 表示可见 return torch.tril(torch.ones(L, L, dtypetorch.bool, devicedevice)) # 组合先算 pad mask再和因果 mask 做广播与运算 src_mask make_pad_mask(src) # (B, 1, Ls) tgt_pad make_pad_mask(tgt_in) # (B, 1, Lt) causal make_causal_mask(tgt_in.size(1), tgt_in.device).unsqueeze(0) # (1, Lt, Lt) tgt_mask tgt_pad causal # (B, Lt, Lt)这里有个形状上的坑tgt_pad是(B, 1, Lt)causal是(1, Lt, Lt)两者做之后广播成(B, Lt, Lt)。广播的结果正好是注意力分数的形状可以直接用。关于因果mask还有个细节它是在每个batch里都是同一份所以完全可以缓存起来不必每次forward都重新生成。序列长度固定的话这个缓存能省下一点开销在长序列上比较可观。提示屏蔽用的是负无穷而不是0。如果你用0去填充被屏蔽的位置softmax之后那些位置仍然会分到正权重只是权重小一点而已pad信息照样会漏进来。这个错误不会报错只会让模型悄悄地学会利用pad。3.3 loss计算ignore_index和log_softmax的关系输出侧的最后一环是损失函数。模型的最后一层是一个线性层把d_model维投影回词表大小得到每个位置的logits。注意这里是logits不是概率不要手贱加softmax——因为CrossEntropyLoss内部已经包含了log_softmax你加了就重复了梯度会异常。class Generator(nn.Module): def __init__(self, d_model, vocab_size): super().__init__() self.proj nn.Linear(d_model, vocab_size) def forward(self, x): # x: (B, L, d_model) - (B, L, V) return self.proj(x)loss计算时需要把(B, L, V)拉平成(B*L, V)标签拉平成(B*L)这是CrossEntropyLoss期望的形状也就是NLLLoss的形状约定。关键是ignore_index参数要设成PAD的id这样被padding的位置不产生损失也不会贡献梯度。criterion nn.CrossEntropyLoss(ignore_index0, label_smoothing0.1) logits model(src, tgt_in, src_mask, tgt_mask) # (B, Lt, V) loss criterion( logits.reshape(-1, logits.size(-1)), # (B*Lt, V) tgt_out.reshape(-1) # (B*Lt,) )label_smoothing这个参数我建议加上尤其是小数据集上。它的作用是把硬标签从1和0软化成0.9和0.1/(V-1)左右防止模型对某一个类别过度自信泛化性会好一些。代价是训练loss数值会变大因为熵本身就高了所以不要拿加了平滑的loss和没加的对比会误判。还有一个常见疑问为什么标签要用tgt_out而不是tgt_in回到3.1那节因为第t个位置的输出预测的是下一个token所以标签必须右移一位。你可以在代码里打印前几条样本确认一下肉眼对一遍比反复读论文快。4. 完整实现用反转任务把闭环跑通前面拆的都是零件这一节把它们拼起来。我用序列反转当验证任务投入最小、反馈最快。整个模型包括三个部分编码器层、解码器层、以及前面讲的输入输出模块。4.1 多头注意力与编码解码层先写多头注意力。这里的核心是把(B, L, d_model)拆成(B, h, L, d_k)。用view加transpose是最省事的写法但要注意view之前张量必须是连续的transpose之后就不连续了所以反向reshape回原形状时要先contiguous()。class MHA(nn.Module): def __init__(self, d_model, n_head, dropout0.1): super().__init__() assert d_model % n_head 0 self.h n_head self.d_k d_model // n_head 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.w_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, q, k, v, maskNone): B, Lq q.size(0), q.size(1) Lk k.size(1) q self.w_q(q).view(B, Lq, self.h, self.d_k).transpose(1, 2) k self.w_k(k).view(B, Lk, self.h, self.d_k).transpose(1, 2) v self.w_v(v).view(B, Lk, self.h, self.d_k).transpose(1, 2) scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(~mask, float(-inf)) attn torch.softmax(scores, dim-1) attn self.dropout(attn) out torch.matmul(attn, v) # (B, h, Lq, d_k) out out.transpose(1, 2).contiguous().view(B, Lq, -1) return self.w_o(out)编码器和解码器层就是注意力加前馈网络每个子层外面套一层残差和LayerNorm。这里我用的是最传统的Post-LN先残差后Norm对新手更直观。如果你要训练深层模型建议换成Pre-LN因为Pre-LN不需要在残差分支上加warmup训练稳定性更好。class FFN(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.net nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) def forward(self, x): return self.net(x) class EncoderLayer(nn.Module): def __init__(self, d_model, n_head, d_ff, dropout0.1): super().__init__() self.attn MHA(d_model, n_head, dropout) self.ffn FFN(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.drop nn.Dropout(dropout) def forward(self, x, mask): x self.norm1(x self.drop(self.attn(x, x, x, mask))) x self.norm2(x self.drop(self.ffn(x))) return x class DecoderLayer(nn.Module): def __init__(self, d_model, n_head, d_ff, dropout0.1): super().__init__() self.self_attn MHA(d_model, n_head, dropout) self.cross_attn MHA(d_model, n_head, dropout) self.ffn FFN(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.drop nn.Dropout(dropout) def forward(self, x, memory, tgt_mask, mem_mask): x self.norm1(x self.drop(self.self_attn(x, x, x, tgt_mask))) x self.norm2(x self.drop(self.cross_attn(x, memory, memory, mem_mask))) x self.norm3(x self.drop(self.ffn(x))) return x交叉注意力那一层的mask要注意它屏蔽的是key和value所在的源序列长度所以形状必须是(B, 1, Ls)而不是目标长度。我第一次写的时候顺手把tgt_mask传进去了seq长度一样时不会报错长度不一样时直接形状不匹配或者更糟——广播成功但屏蔽了错误的位置。4.2 数据构造与训练循环反转任务的构造很简单随机生成一串token目标是它的逆序。这里我用torch.randint直接生成不需要任何外部数据文件方便你复制过去就运行。PAD, BOS, EOS, VOCAB_SIZE 0, 1, 2, 50 def make_batch(batch_size32, seq_len8): # 生成范围 [3, VOCAB_SIZE) 的token避开特殊符号 src torch.randint(3, VOCAB_SIZE, (batch_size, seq_len)) rev torch.flip(src, dims[1]) bos torch.full((batch_size, 1), BOS) eos torch.full((batch_size, 1), EOS) tgt torch.cat([bos, rev, eos], dim1) # [BOS, rev..., EOS] return src, tgt训练循环里我用了Noam式学习率也就是线性warmup加逆平方根衰减。它对Transformer这种深层结构几乎是标配因为训练初期梯度方向很不稳定学习率太大会直接发散。这个调度器不需要手调warmup_steps设成总步数的十分之一左右就够了。def train(steps800, d_model64, n_head4, n_layer2, d_ff128): model Seq2Seq(d_model, n_head, n_layer, d_ff, VOCAB_SIZE).to(DEVICE) opt torch.optim.Adam(model.parameters(), lr0.0, betas(0.9, 0.98), eps1e-9) criterion nn.CrossEntropyLoss(ignore_indexPAD, label_smoothing0.0) for step in range(1, steps 1): # warmup 逆平方根衰减 lr d_model ** -0.5 * min(step ** -0.5, step * 400 ** -1.5) for g in opt.param_groups: g[lr] lr src, tgt make_batch() src, tgt src.to(DEVICE), tgt.to(DEVICE) tgt_in, tgt_out tgt[:, :-1], tgt[:, 1:] src_mask make_pad_mask(src) tgt_mask make_pad_mask(tgt_in) make_causal_mask(tgt_in.size(1), DEVICE).unsqueeze(0) logits model(src, tgt_in, src_mask, tgt_mask) loss criterion(logits.reshape(-1, logits.size(-1)), tgt_out.reshape(-1)) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() if step % 100 0: print(fstep {step:4d} | lr {lr:.2e} | loss {loss.item():.4f})梯度裁剪那行不要省。Transformer的梯度在训练初期会非常激进尤其是最后的投影层偶尔会冒出一个很大的梯度把参数带偏。裁剪阈值1.0是经验值一般不用改。如果发现loss偶尔出现尖刺然后又恢复八成就是没裁剪。跑起来之后你应该看到这样的曲线前一百步loss从4左右降到0.5以下三百步左右跌到接近0因为词表只有50、序列长度只有8任务非常容易。如果到了五百步loss还在1.5以上那基本可以确定是mask或者标签那里出了错不要怀疑模型容量去检查shift。4.3 推理阶段自回归生成训练时的输入输出都是已知的推理时不是。推理要从一个BOS开始一步步生成每生成一个token就把它拼接到输入末尾再让模型预测下一个。这个过程叫自回归解码。torch.no_grad() def greedy_decode(model, src, max_len20): model.eval() src_mask make_pad_mask(src) memory model.encode(src, src_mask) B src.size(0) ys torch.full((B, 1), BOS, dtypetorch.long, devicesrc.device) finished torch.zeros(B, dtypetorch.bool, devicesrc.device) for _ in range(max_len): tgt_mask make_pad_mask(ys) make_causal_mask(ys.size(1), ys.device).unsqueeze(0) out model.decode(memory, ys, tgt_mask, src_mask) logits model.generator(out[:, -1]) # 只取最后一个位置 next_id logits.argmax(dim-1, keepdimTrue) # (B, 1) ys torch.cat([ys, next_id], dim1) finished | (next_id.squeeze(1) EOS) if finished.all(): break return ys这段代码有几个点值得强调。第一model.eval()和torch.no_grad()必须加前者关掉dropout后者省显存。第二每次循环都要重新计算tgt_mask因为序列长度在变因果mask的尺寸也跟着变。第三只需要取最后一个位置的logits因为前面的位置早就在上一步生成过了重复计算是浪费——这也是KV缓存要解决的问题。第四EOS判断是逐样本的一个batch里有的样本先结束理论上可以给它补PAD并跳过计算但为了代码简单我先没做。如果你要上生产这里至少有两个优化空间。一是KV缓存把每层注意力的key和value缓存下来每一步只算新token复杂度从O(L²)降到O(L)。二是批量采样策略贪心解码容易生成重复内容换成beam search或者top-k采样会自然很多。这两块展开又是一篇这里先提一句知道方向就行。5. 排查实录我踩过的坑和对应解法这部分是全篇最实用的。前面讲的是应该怎么做这里讲的是做错了会怎样、怎么找回来。我把遇到的问题按症状分类整理你可以当速查表用。5.1 形状与类型类问题速查表形状错误是最好查的因为PyTorch会直接报错。麻烦的是那些不报错但结果悄悄错的我单独在后面列。报错/症状大概率原因定位方法matmul维度不匹配mask形状与scores不匹配在masked_fill前打印mask和scores的shapemasked_fill dtype错误mask用了float不是bool统一用(seq ! PAD)生成得到bool位置编码相加报错pe没注册成buffer还在CPU检查register_buffer是否正确调用传给交叉注意力的mask长度不对用了tgt_mask而非src_mask交叉注意力的mask长度必须等于源序列长度CrossEntropyLoss维度错误logits没reshape成二维用.reshape(-1, V)压平推理时越界max_len超过了位置编码表长度检查max_len设置我这个表格里最常出现的是第二条。因为很多老教程用的是加性mask0和负无穷生成的mask是float类型masked_fill要求布尔两者混用就会报错。我的建议是全程只用布尔mask加性mask如果需要比如实现相对位置偏置再单独处理。5.2 训练不收敛的几种典型表现第一种loss从第一步就在4.0附近不动或者直接是nan。这种情况先查学习率warmup没做好或者初始lr太大是最常见的。其次查标签里有没有出现超出词表范围的id——CrossEntropyLoss在标签越界时会直接给nan而不是报错非常坑。我养成的习惯是每次造完数据加一句断言assert tgt.max() vocab_size。第二种loss能降到2左右但卡住不动。这个位置恰好对应均匀分布的熵说明模型学到的就是词频分布没有任何结构信息。原因通常是mask写反了把该保留的位置屏蔽了、该屏蔽的位置保留了。特别是因果mask如果上下三角搞反模型只能看到未来的token当然学不出东西。验证方法很简单把mask矩阵打印出来看下三角是不是True。第三种训练loss很低但验证集或者推理结果很差。这就是典型的过拟合或训练推理不一致。如果不一致去检查三点推理时有没有加因果mask、有没有误用dropout、EOS有没有正确终止。我遇到过一次是推理时忘了把模型切到eval模式dropout还在生效输出全是乱的。第四种loss正常但在某个值附近震荡不下降。这时候把学习率调小十倍试试如果还是不行检查一下label_smoothing——加了平滑之后loss有下限别把它当成不收敛。提示每次调参只改一个变量。我早期吃过这个亏一次改三处效果变好了也不知道是哪一处起的作用下次遇到问题还是不知道怎么复现。5.3 几个我希望早点知道的经验关于padding_idx前面说过它只保护Embedding自己。但还有个隐藏行为如果你在位置编码里对pad位置也加了位置向量那pad位置就不再是全零了它会变成一个纯位置向量依然会污染注意力。解决办法很简单做完位置编码后把pad位置乘回一个掩码x x.masked_fill((tgt_in PAD).unsqueeze(-1), 0.0)关于batch_firstPyTorch原生的nn.Transformer默认是batch_firstFalse也就是序列维在前。这是历史遗留问题很多人包括我第一次用的时候根本没注意模型跑起来了但结果全错。现在的版本支持传batch_firstTrue我强烈建议打开它整条链路都用(B, L, d)的约定不用来回转置。手写的话就更简单了一开始就坚持(B, L, d)别给自己添麻烦。关于显存注意力矩阵是O(L²)的序列长度从512涨到1024显存直接翻四倍。手撕版本在调试阶段很容易OOM我的做法是先把batch和序列长度压到很小比如批量8、长度16把逻辑跑通再加规模。别一上来就上真实数据集的规模调不完的。关于可复现性训练前一定固定随机种子包括torch.manual_seed、random.seed、numpy.random.seed并且注意CuDNN的确定性选项。我吃过一次亏同一份代码跑三遍得到三个不同的收敛速度花了半天时间怀疑人生最后发现是随机种子没固定。这个习惯在小规模调试阶段尤其重要不然你根本分不清是改代码起了作用还是运气好。最后说个心态上的事。手撕Transformer的意义不在于以后真的手写而在于你能读懂报错的含义、能判断一个开源实现哪里不对劲、能在模型不工作时快速定位到是数据、mask还是loss的问题。我见过太多人调包调得很顺一遇到形状不匹配就束手无策。多写两遍把这条链路刻进肌肉记忆里后面看任何变体——无论它是改注意力、改位置编码还是改归一化位置——你心里都会有一个稳定的坐标系。
