很多人一开始接触Transformer就是调库像我最早也是直接拿现成工具包做情感分类跑通了就觉得已经会了。后来真正从零写了一个MiniTransformer才意识到框架把太多细节藏在了黑盒底下——注意力分数为什么除以根号d层归一化放哪里为什么会影响训练稳定性位置编码选型为什么直接决定外推能力这些问题在文档里根本找不到答案。这篇是Transformer自然语言处理系列第三篇我索性把手写解码器模型的完整过程整理出来从模型骨架到分词训练再到生成和部署尽量把这几年踩过的坑一次说清楚。适合那些已经熟悉PyTorch基础操作、但想真正把Transformer内部结构摸透的读者哪怕你暂时不打算从零训练大模型搞懂这些对调试别人的代码也很有帮助。1. 为什么第三篇决定手写Transformer框架黑盒的代价1.1 框架用久了反而不知道模型在干什么我见过不少同学用开源代码微调大模型头头是道Loss曲线调得一帆风顺但一旦模型崩溃就完全无从下手。原因很简单他们对模型内部发生的计算没有具体感知。手写一个训练规模不大的Transformer从零开始的价值就在这里——你能亲眼看到每一步张量的形状怎么变化注意力矩阵长什么样梯度在反向传播时数值是否异常。这个可视化心智模型一旦建立起来之后再去读那些工业级代码库速度会快一个量级。更实际的好处是调试。我自己写的时候曾经因为注意力掩码矩阵的维度排列问题让模型悄悄看到了未来时刻的信息训练出来的loss漂亮得离谱生成的时候却驴唇不对马嘴。这种错误在调库的场景下几乎不可能定位但在手写代码里我逐步检查张量shape就能一眼看出问题。1.2 手写的目标只做最核心的Decoder-only我们从零搭建Transformer通常有三条路线完整的Encoder-Decoder结构适合翻译任务Encoder-only适合理解任务Decoder-only适合生成任务。本文选Decoder-only原因有两点。第一现在主流自然语言处理应用大都围绕生成展开Decoder-only结构不仅能做续写加一些改造就能做分类实际应用价值最大。第二它的掩码机制——因果掩码——让实现过程更能暴露Transformer的细节难点。你一旦把Decoder-only从头到尾写通了再回头理解编码器的自注意力简直像看简化版。定了结构之后就是范围控制。我的目标是一个参数规模在150M以内的字符级/BPE混合小模型单卡可以训练。别一上来就想复刻几百B参数的大模型先把推理链路做通把注意力机制写透比追求规模有价值得多。2. 手写Transformer从多头注意力到位置编码的取舍2.1 模型维度怎么定d_model、num_heads与head_dim手写模型的第一个决策点就是维度设计。我先给出一套我自己用的参数配置后面所有代码都基于这套配置展开。参数数值说明vocab_size32000取决于BPE词典大小实际训练前确定d_model512词向量及隐藏层维度num_heads8注意力头数head_dim64等于d_model / num_heads这是硬性要求num_layers8解码器层数max_seq_len512训练序列最大长度dropout0.1训练时默认dropout比例这里关键的一条是head_dim必须等于d_model整除num_heads的结果。多头注意力的做法是把d_model维度的向量切分成num_heads份每一份做独立的自注意力再拼接回去。如果d_model512num_heads8head_dim就是64但是如果d_model512num_heads7那就没法均匀切分实现上还得特判padding非常麻烦所以设计模型时尽量让除法整除。每个头的计算本身和单头注意力没本质区别。对第i个头Q、K、V分别经过独立的线性投影得到形状为(batch_size, seq_len, head_dim)的张量然后计算缩放点积注意力import torch import torch.nn.functional as F def scaled_dot_product_attention(q, k, v, maskNone, dropoutNone): d_k q.size(-1) # q, k, v: (batch, heads, seq_len, head_dim) scores q k.transpose(-2, -1) / (d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) weights torch.softmax(scores, dim-1) if dropout is not None: weights dropout(weights) return weights v, weights为什么除以根号d这个细节值得说。如果不缩放当head_dim较大时点积结果方差也会变大softmax会掉进梯度极小的饱和区反向传播时梯度接近于零训练会非常慢甚至完全停摆。除以根号d可以保证点积结果的方差稳定在一个合适量级这个操作是整套Transformer训练的基石之一。2.2 因果掩码为什么mask矩阵必须做成下三角Decoder-only模型里的掩码稍微复杂一点。核心逻辑是自助回归生成时第t个位置只能看到第t个位置及之前的信息不能偷看后面的词。实现上就是把softmax前的score矩阵中右上角部分填充为负无穷def build_causal_mask(seq_len): return torch.tril(torch.ones(seq_len, seq_len)).bool()配合批量训练时还需要处理padding部分。一个batch里的序列长度各不相同短序列的尾部有padding token这部分一定不能参与注意力计算。处理方法是在attention里同时传入两个掩码一个是causal mask一个是padding maskpadding mask用来把padding位置的score也置为无穷小。我第一次写就栽在这两套mask的合并上。正确做法是causal mask和padding mask做逻辑与只保留既属于因果可见范围、又不是padding位置的格子。代码里用bool masked_fill一次搞定# (batch, 1, seq_len, seq_len) combined_mask causal_mask padding_mask.unsqueeze(1).unsqueeze(2) scores scores.masked_fill(~combined_mask, float(-inf))2.3 位置编码选型可学习位置编码与旋转位置编码Transformer本身没有顺序概念引入位置信息的方案有好几种。初版论文用正弦余弦编码后来的实践分化成两个方向可学习位置编码和旋转位置编码RoPE。可学习位置编码就是把位置向量做成一个可训练的嵌入表简单直接。它的优点是实现容易在小模型上效果也很稳缺点是外推能力差训练时最大长度512到推理时一旦输入长度超过512模型表现会急剧下降。旋转位置编码则更优雅它把相对位置信息通过旋转矩阵注入到Q、K向量中核心性质是内积只与相对位置有关。RoPE在长文本场景下明显更占优但是实现复杂度也更高。我做这个项目时选择的是可学习位置编码。原因太简单训练语料长度普遍不超过512短文本上两种方案差异几乎可以忽略而可学习编码实现成本低、更容易把注意力训练稳定。如果后期确实要处理长文本切换到RoPE也不难训练代码主体不用大改。2.4 Layer Norm放哪里Pre-LN和Post-LN的分水岭初版Transformer把LayerNorm放在每个子层之后这种叫Post-LN现代很多实现把LayerNorm放在子层输入处这叫Pre-LN。两者看起来只是挪了个位置实际训练稳定性差距很大。Post-LN在训练初期特别容易炸因为梯度要穿过残差连接和LayerNorm叠加的路径数值稍有不慎就发散。Pre-LN把LayerNorm放在残差连接的起点等于给每一条子层路径都做了标准化梯度流更平稳所以现在大部分开源模型都用Pre-LN。我选择Pre-LN还用了一个关键初始化技巧对每个子层最后的线性输出乘上一个缩放系数系数与层数相关常见做法是乘上(2 * num_layers)^-0.5。这个操作能在网络较深时抑制梯度方差爆炸。用上之后8层模型训练稳定度肉眼可见地提升。2.5 前馈网络从ReLU到GELU再到GLUTransformer每个解码器层里除了注意力还有一层位置前馈网络。经典版本是两层线性ReLU激活高维映射d_model从512升到2048再降回512。现代实现大多用GELU激活因为它的曲线更平滑梯度不像ReLU那样在零点硬变向。我自己的经验是PyTorch里直接F.gelu就好无缝替换。更激进的选择是GLU变体它在某些场景下能提升表现但参数量和计算量会增加。对这个小模型够用就好。class FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.w1 nn.Linear(d_model, d_ff) self.w2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) def forward(self, x): return self.w2(self.dropout(F.gelu(self.w1(x))))3. 让模型真正开始说人话数据、分词与训练策略3.1 训练语料的预处理思路有了模型结构接下来的一切都围绕数据展开。我在这个项目中使用的语料是一批公开爬取的中文与英文混合文本数量级在几千万行。数据清洗这一步看似不起眼实际决定训练成败。清洗流程按优先级排先去除HTML标签和无意义符号再过滤掉过短和过长的行接着做段落切分最后按连续性拼接成固定大小的训练切片。这些操作看似和Transformer无关但脏数据会让分词器学到一堆垃圾token模型后续生成质量直接打折。还要特别注意语料的平衡性。如果全是短句模型只能学会输出短句如果混入了大量文本稀疏但显而易见过拟合的样板文本模型生成时会频繁复读这些套路句子。行业里管这叫数据污染效果非常隐蔽初期训练loss正常生成却崩得一塌糊涂很多新人排查半天都找不到原因。3.2 分词器选型字符级还是BPE小规模训练时有两个选择字符级分词和BPE分词。字符级就是把每个汉字、每个英文字母当成一个token词表大小几百实现省事但序列会被拉得很长模型学语义的表达效率低。BPE则是从数据中统计学出子词切分规则比如自然语言处理可能被切成自然、语言、处理三段子词词表压缩效率高。我的建议是如果语料总量在1亿字符以上直接上BPE如果只是做小玩具验证字符级也能跑通。我在这里用的是SentencePiece的BPE模式词表32000加上特殊tokenpad、bos、eos、unk。分词器训练也要匹配推理时场景千万别用训练语料之外的分词器否则token id对不上模型输出就是乱码。这个低级错误我踩过不止一次。3.3 数据加载动态padding比固定长度省资源混合语料里文本长度天然长短不一。常见做法是统一截断成固定长度用固定长度batch训练简单但浪费。我的做法是动态padding按文本长度排序分组每次组内取最大长度为batch尺寸padding到这个最大值整体显存利用率和训练速度都有提升。实际代码在torch.utils.data.DataLoader里加一个collate_fn就行。关键点是训练阶段对padding部分要加掩码前面已经提过预测阶段则不需要padding mask因为是单条推理。还有一个细节是采样顺序。如果完全随机采样相邻样本长短差异过大动态padding的收益就不明显。我用的是bucket sampling策略按长度分桶桶内随机打乱这样每个batch内的序列长度接近padding率能降到20%以下。3.4 优化器与学习率调度AdamW与WarmupTransformer训练的优化器选择基本没什么悬念AdamW。AdamW和传统Adam的区别在于权重衰减的实现方式——AdamW把权重衰减从梯度更新中解耦能有效降低过拟合风险现在已经是主流默认。学习率调度上我强烈建议使用warmup配合后续衰减。Warmup阶段学习率从极小值线性爬升到峰值作用是让模型在前几百步先稳住注意力头和LayerNorm的参数避免一开始就用大学习率把梯度更新方向带偏。峰值后可以用cosine退火慢慢降低学习率直到训练结束。实际参数峰值学习率5e-4warmup步数占训练总步数的2%左右比如总步数20000warmup就400步。这个比例是我从多次实验中总结出来的经验值太大会延长收敛时间太小则容易在初期就出现loss抖动。固定好这些参数之后我的Loss曲线基本是一条平滑下降的弧线。3.5 损失函数CrossEntropy与标签平滑生成任务的损失就是下一步token预测的交叉熵。这里有一个容易被忽略的细节padding位置不参与loss计算。做法是把logits和label都压平后只取label不是pad的部分做cross entropy。标签平滑是另一个值得开的开关。它会把唯一正确答案的一-hot标签稍微抹平让模型不至于过度自信对缓解过拟合和提升泛化能力都有帮助。这里设置为0.1即在损失计算时把真实token的置信度从1降低到0.9剩下的0.1平均分配给其他词表token。顺带提醒一点如果发现模型训练时loss总是卡在自然语言模型的理论下界附近甚至低于论文报告的数字先别高兴很可能是学习率设置太低或者标签平滑失效导致了某种梯度坍缩。4. 训练MiniTransformer时最常翻车的三个点4.1 损失原地不动或稳步走高先查学习率再查掩码训练开始后如果发现Loss前一百步纹丝不动我第一个怀疑的是学习率。太大或太小都会让模型停在原地太大时甚至直接NaN。调试方法是用一个极小的子集跑几十步逐渐升高学习率观察Loss是否出现先降后升的曲线由此找到合理区间。如果学习率正常但Loss依然不动就要手动验证注意力掩码。我前面提到的因果掩码错误曾让我白白烧了大半天训练时间。怎么自查写一段很小的序列模拟前向计算存下attention权重人工检查第t行是否只有前t列有非零值。只要有一步越界立刻就能看出来。还有一个隐蔽的坑是embedding层和最后的输出层共享权重时词表过大可能会导致输出层梯度爆炸。解决办法是输出层不要直接复用embedding权重或者单独加一层LayerNorm做输出标准化。这两种做法在小模型上都验证有效。4.2 显存报错和梯度截断训练时经常遇到的第二个大坑是显存问题。8层Decoder batch size 32 seq_len 512在单张消费级显卡上很容易吃紧。我的处理顺序是先开PyTorch的混合精度训练也就是AMP显存能省一半左右速度也能提升还不够就开gradient checkpointing在每一层前向时丢弃中间激活反向传播时重新计算用时间换空间最后还是不够再调小batch。梯度截断这个操作最好养成习惯。torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)放在loss.backward()之后、optimizer.step()之前。它能在某些异常batch产生巨大梯度时避免整个模型的参数被彻底冲坏训练稳定性的提升非常显著。4.3 过拟合、记忆化与幻觉问题往往不在模型结构里损失一直降得很好但生成出来的文本反复出现训练语料里的原文段落这是典型的记忆化现象。说明小模型没有真正学会语言规律只是在背诵语料片段。缓解办法有几种提升训练数据多样性、加大dropout、使用标签平滑、降低模型层数或宽度。更深层的原因是模型参数规模相对语料太小或太大时都会诱导记忆化。我亲身实验过把模型从150M缩到40M时反而出现了更多复读原文的现象因为规模不够大根本没法把复杂语法规律抽象出来只能背。这个度需要靠实验来确定不要一味堆参数也不要认为模型越小越安全。5. 推理解码生成质量与速度的平衡5.1 采样策略对比贪心、Beam Search与Top-p模型训练完成之后生成文本的方式对结果影响特别大。贪心解码每次取概率最大的token速度快但容易陷入循环且缺少多样性。Beam Search在机器翻译上表现好但在开放性生成上经常产生重复严重的句子。Top-k和Top-p是目前的主流选择。Top-k限制候选集为概率最高的k个tokenTop-p则按累计概率动态截断候选集。两者可以结合使用我在小模型上最常用的配置是k50同时p0.9能比较好地平衡多样性和流畅度。采样还受到temperature控制。temperature大于1时分布变平输出更加随机小于1时分布变尖输出更保守。实际生成时我一般从temperature0.8开始调然后再微调k和p每次只看十几条输出就能判断方向。5.2 重复生成病灶n-gram惩罚与重复惩罚系数生成质量最典型的问题是复读机。单靠采样策略很难根除通常需要配合重复惩罚。最粗暴有效的方式是no_repeat_ngram_size3强制模型不生成任何重复的三元组。这种硬性限制能消灭大部分循环但有可能切断一些本来合理的表达长文本生成到最后会变得生硬。另一个手段是Fisher惩罚系数的变体即在候选token概率上乘以下面这个值出现次数越多惩罚越大。我把系数设在1.2到1.5之间它能有效减重复同时又不会像ngram惩罚那样完全禁止可能出现的长序列重复结构。我建议的组合是短文本生成时只给温度Top-p不加强制惩罚长文本生成时再加no_repeat_ngram_size和重复惩罚。两个场景分开调试比单一配置通打天下效果更好。5.3 推理速度优化KV Cache和批处理生成阶段的慢主要来自每次只解码一个token但又要对整段历史做完整前向。最成熟的优化就是KV Cache把之前所有位的Key和Value缓存下来新token只需计算自身的Q再与缓存拼接后的K、V做注意力。实现上就是在解码循环里维护两个缓存张量把每一步新算出的K、V沿seq_len维度concat进去。第一次解码时缓存为空后面的步数全部复用。实测KV Cache在512长度序列上解码速度能提升3-5倍。如果还想更快可以把多个请求拼成一个batch一起做推理GPU吞吐会大幅上涨但要注意各条序列的padding处理。6. 从玩具到实际工程微调、量化和部署的取舍6.1 从零训练后要不要在开源基座上做微调手写训练出来的小模型生成质量和通用能力肯定是无法和成熟的开源模型比的。如果想要得到一个可用性更强的小模型最可行的路线是拿着别人训练好的基座模型做下游任务的微调而不是从零预训练。微调阶段要关注三个点下游任务数据的质量是第一优先数据条数从几千到几万都有作用关键是领域覆盖要准学习率要比预训练低一到两个数量级比如5e-5到1e-4冻结部分底层参数只微调高层也可以显著降低过拟合风险。这些经验在我做具体下游任务时反复验证过。6.2 轻量化部署的路子量化与导出模型训练好不等于能上线。部署时先把模型导出成推理引擎格式比如ONNX再开启动态量化把权重从FP32压到INT8模型体积和推理延迟都会明显下降。量化的代价是精度损失尤其在计算密集的注意力部分实测在分类任务上F1下降通常在1%以内文本生成任务上则可能更明显。如果显存还紧张还可以考虑蒸馏。用一个更大的教师模型指导小模型学习让输出概率分布尽量接近教师模型。蒸馏后的小模型在保持较小体积的同时效果往往比直接训练同样规模的新模型要好。这条路在移动端或低配CPU场景几乎是标配。6.3 这次手写之后我对Transformer的重新理解最后说点个人体会。这个项目做完之后我最大的变化是再也敢只看图就能说懂Transformer了而是习惯性地从具体张量形状怎么流动去理解每一篇新模型结构。以后遇到open-source代码里一个奇怪的reshape或者transpose我能很快意识到它服务于哪些约束。另一个收获是训练和调试的依赖关系被彻底打通。以前遇到Loss曲线异常只会换数据换学习率现在会先去看模型的归一化层分布、注意力权重、梯度范数这些内部信号定位问题的速度比之前快了许多。这次代码整体规模不大但足以成为理解现代语言模型的支点。如果你也想把Transformer彻底吃透我很建议你也动手写一遍哪怕只是实现一个能跑通的最简版本回报率远超想象。